Repository: stanford-oval/storm Branch: main Commit: fb951af7744d Files: 66 Total size: 581.6 KB Directory structure: gitextract_m_9bgqy0/ ├── .github/ │ ├── ISSUE_TEMPLATE/ │ │ └── bug_report.md │ └── workflows/ │ ├── format-check.yml │ └── python-package.yml ├── .gitignore ├── .pre-commit-config.yaml ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── examples/ │ ├── costorm_examples/ │ │ └── run_costorm_gpt.py │ └── storm_examples/ │ ├── README.md │ ├── helper/ │ │ └── process_kaggle_arxiv_abstract_dataset.py │ ├── run_storm_wiki_claude.py │ ├── run_storm_wiki_deepseek.py │ ├── run_storm_wiki_gemini.py │ ├── run_storm_wiki_gpt.py │ ├── run_storm_wiki_gpt_with_VectorRM.py │ ├── run_storm_wiki_groq.py │ ├── run_storm_wiki_mistral.py │ ├── run_storm_wiki_ollama.py │ ├── run_storm_wiki_ollama_with_searxng.py │ └── run_storm_wiki_serper.py ├── frontend/ │ └── demo_light/ │ ├── .streamlit/ │ │ └── config.toml │ ├── README.md │ ├── demo_util.py │ ├── pages_util/ │ │ ├── CreateNewArticle.py │ │ └── MyArticles.py │ ├── requirements.txt │ ├── stoc.py │ └── storm.py ├── knowledge_storm/ │ ├── __init__.py │ ├── collaborative_storm/ │ │ ├── __init__.py │ │ ├── engine.py │ │ └── modules/ │ │ ├── __init__.py │ │ ├── article_generation.py │ │ ├── callback.py │ │ ├── co_storm_agents.py │ │ ├── collaborative_storm_utils.py │ │ ├── costorm_expert_utterance_generator.py │ │ ├── expert_generation.py │ │ ├── grounded_question_answering.py │ │ ├── grounded_question_generation.py │ │ ├── information_insertion_module.py │ │ ├── knowledge_base_summary.py │ │ ├── simulate_user.py │ │ └── warmstart_hierarchical_chat.py │ ├── dataclass.py │ ├── encoder.py │ ├── interface.py │ ├── lm.py │ ├── logging_wrapper.py │ ├── rm.py │ ├── storm_wiki/ │ │ ├── __init__.py │ │ ├── engine.py │ │ └── modules/ │ │ ├── __init__.py │ │ ├── article_generation.py │ │ ├── article_polish.py │ │ ├── callback.py │ │ ├── knowledge_curation.py │ │ ├── outline_generation.py │ │ ├── persona_generator.py │ │ ├── retriever.py │ │ └── storm_dataclass.py │ └── utils.py ├── requirements.txt └── setup.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .github/ISSUE_TEMPLATE/bug_report.md ================================================ --- name: Bug report about: Create a report to help us improve title: "[BUG]" labels: '' assignees: '' --- **Describe the bug** A clear and concise description of what the bug is. **To Reproduce** Report following things 1. Input topic name 2. All output files generated for this topic as a zip file. **Screenshots** If applicable, add screenshots to help explain your problem. **Environment:** - OS: [e.g. iOS, Windows] - Browser [e.g. chrome, safari] if the bug report is UI problem ================================================ FILE: .github/workflows/format-check.yml ================================================ name: Check Python formatting with Black on: pull_request: branches: - main jobs: lint: runs-on: ubuntu-latest steps: - uses: actions/checkout@v2 - uses: actions/setup-python@v2 - uses: psf/black@stable with: black_args: "knowledge_storm --check" ================================================ FILE: .github/workflows/python-package.yml ================================================ name: Build and upload Python package on: workflow_dispatch: # Allows manual triggering of the workflow jobs: build: runs-on: ubuntu-latest steps: - uses: actions/checkout@master - name: Set up Python 3.11 uses: actions/setup-python@v3 with: python-version: "3.11" - name: Compare versions in setup.py and knowledge_storm/__init__.py run: | VERSION_SETUP=$(grep -oP '(?<=version=\").*(?=\")' setup.py) VERSION_INIT=$(grep -oP '(?<=__version__ = \").*(?=\")' knowledge_storm/__init__.py) echo "Version in setup.py: $VERSION_SETUP" echo "Version in __init__.py: $VERSION_INIT" if [ "$VERSION_SETUP" != "$VERSION_INIT" ]; then echo "Error: Version mismatch between setup.py ($VERSION_SETUP) and knowledge_storm/__init__.py ($VERSION_INIT)" exit 1 fi shell: bash - name: Install dependencies run: python3 -m pip install setuptools wheel twine - name: Install dependencies run: | python3 -m pip install --upgrade pip setuptools wheel if [ -f requirements.txt ]; then pip install -r requirements.txt; fi - name: Build a binary wheel run: python3 setup.py sdist bdist_wheel - name: Publish package to PyPI uses: pypa/gh-action-pypi-publish@release/v1 with: user: __token__ password: ${{ secrets.PYPI_API_TOKEN }} ================================================ FILE: .gitignore ================================================ # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # mac .DS_Store # Other .vscode *.tsv *.pt gpt*.txt *.env local/ local_* build/ *.egg-info/ .idea .venv # Project-specific secrets.toml *.log */assertion.log *results/ .venv/ ================================================ FILE: .pre-commit-config.yaml ================================================ repos: - repo: https://github.com/psf/black rev: 24.8.0 hooks: - id: black name: Format Python code with black entry: black args: ["knowledge_storm/"] language: python pass_filenames: true ================================================ FILE: CONTRIBUTING.md ================================================ # Contributing Thank you for your interest in contributing to STORM! Contributions aren't just about code. Currently (last edit: 7/22/2024), we are accepting the following forms of contribution: - Pull requests for additional language model support to `knowledge_storm/lm.py`. - Pull requests for additional retrieval model/search engine support to `knowledge_storm/rm.py`. - Pull requests for new features to `frontend/demo_light` to assist other developers. - Identification and reporting of issues or bugs. - Helping each other by responding to issues. Please note that we are not accepting code refactoring PRs at this time to avoid conflicts with our team's efforts. ## Development This section contains technical instructions & hints for contributors. ### Setting up 1. Fork this repository and clone your forked repository. 2. Install the required packages: ``` conda create -n storm python=3.11 conda activate storm pip install -r requirements.txt ``` 3. If you want to contribute to `frontend/demo_light`, follow its [Setup guide](https://github.com/stanford-oval/storm/tree/main/frontend/demo_light#setup) to install additional packages. ### PR suggestions Following the suggested format can lead to a faster review process. **Title:** [New LM/New RM/Demo Enhancement] xxx **Description:** - For new language model support, (1) describe how to use the new LM class, (2) create an example script following the style of existing example scripts under `examples/`, (3) attach an input-output example of the example script. - For new retrieval model/search engine support, (1) describe how to use the new RM class and (2) attach input-output examples of the RM class. - For demo light enhancements, (1) describe what's new and (2) attach screenshots to demonstrate the UI change. - Please clearly describe the required API keys and provide instructions on how to get them (if applicable). This project manages API key with `secrets.toml`. **Code Format:** We adopt [`black`](https://github.com/psf/black) for arranging and formatting Python code. To streamline the contribution process, we set up a [pre-commit hook](https://pre-commit.com/) to format the code under `knowledge_storm/` before committing. To install the pre-commit hook, run: ``` pip install pre-commit pre-commit install ``` The hook will automatically format the code before each commit. ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2024 Stanford Open Virtual Assistant Lab Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: MANIFEST.in ================================================ include requirements.txt include LICENSE include README.md ================================================ FILE: README.md ================================================

# STORM: Synthesis of Topic Outlines through Retrieval and Multi-perspective Question Asking

| Research preview | STORM Paper| Co-STORM Paper | Website |

**Latest News** 🔥 - [2025/01] We add [litellm](https://github.com/BerriAI/litellm) integration for language models and embedding models in `knowledge-storm` v1.1.0. - [2024/09] Co-STORM codebase is now released and integrated into `knowledge-storm` python package v1.0.0. Run `pip install knowledge-storm --upgrade` to check it out. - [2024/09] We introduce collaborative STORM (Co-STORM) to support human-AI collaborative knowledge curation! [Co-STORM Paper](https://www.arxiv.org/abs/2408.15232) has been accepted to EMNLP 2024 main conference. - [2024/07] You can now install our package with `pip install knowledge-storm`! - [2024/07] We add `VectorRM` to support grounding on user-provided documents, complementing existing support of search engines (`YouRM`, `BingSearch`). (check out [#58](https://github.com/stanford-oval/storm/pull/58)) - [2024/07] We release demo light for developers a minimal user interface built with streamlit framework in Python, handy for local development and demo hosting (checkout [#54](https://github.com/stanford-oval/storm/pull/54)) - [2024/06] We will present STORM at NAACL 2024! Find us at Poster Session 2 on June 17 or check our [presentation material](assets/storm_naacl2024_slides.pdf). - [2024/05] We add Bing Search support in [rm.py](knowledge_storm/rm.py). Test STORM with `GPT-4o` - we now configure the article generation part in our demo using `GPT-4o` model. - [2024/04] We release refactored version of STORM codebase! We define [interface](knowledge_storm/interface.py) for STORM pipeline and reimplement STORM-wiki (check out [`src/storm_wiki`](knowledge_storm/storm_wiki)) to demonstrate how to instantiate the pipeline. We provide API to support customization of different language models and retrieval/search integration. [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) ## Overview [(Try STORM now!)](https://storm.genie.stanford.edu/)

STORM is a LLM system that writes Wikipedia-like articles from scratch based on Internet search. Co-STORM further enhanced its feature by enabling human to collaborative LLM system to support more aligned and preferred information seeking and knowledge curation. While the system cannot produce publication-ready articles that often require a significant number of edits, experienced Wikipedia editors have found it helpful in their pre-writing stage. **More than 70,000 people have tried our [live research preview](https://storm.genie.stanford.edu/). Try it out to see how STORM can help your knowledge exploration journey and please provide feedback to help us improve the system 🙏!** ## How STORM & Co-STORM works ### STORM STORM breaks down generating long articles with citations into two steps: 1. **Pre-writing stage**: The system conducts Internet-based research to collect references and generates an outline. 2. **Writing stage**: The system uses the outline and references to generate the full-length article with citations.

STORM identifies the core of automating the research process as automatically coming up with good questions to ask. Directly prompting the language model to ask questions does not work well. To improve the depth and breadth of the questions, STORM adopts two strategies: 1. **Perspective-Guided Question Asking**: Given the input topic, STORM discovers different perspectives by surveying existing articles from similar topics and uses them to control the question-asking process. 2. **Simulated Conversation**: STORM simulates a conversation between a Wikipedia writer and a topic expert grounded in Internet sources to enable the language model to update its understanding of the topic and ask follow-up questions. ### CO-STORM Co-STORM proposes **a collaborative discourse protocol** which implements a turn management policy to support smooth collaboration among - **Co-STORM LLM experts**: This type of agent generates answers grounded on external knowledge sources and/or raises follow-up questions based on the discourse history. - **Moderator**: This agent generates thought-provoking questions inspired by information discovered by the retriever but not directly used in previous turns. Question generation can also be grounded! - **Human user**: The human user will take the initiative to either (1) observe the discourse to gain deeper understanding of the topic, or (2) actively engage in the conversation by injecting utterances to steer the discussion focus.

Co-STORM also maintains a dynamic updated **mind map**, which organize collected information into a hierarchical concept structure, aiming to **build a shared conceptual space between the human user and the system**. The mind map has been proven to help reduce the mental load when the discourse goes long and in-depth. Both STORM and Co-STORM are implemented in a highly modular way using [dspy](https://github.com/stanfordnlp/dspy). ## Installation To install the knowledge storm library, use `pip install knowledge-storm`. You could also install the source code which allows you to modify the behavior of STORM engine directly. 1. Clone the git repository. ```shell git clone https://github.com/stanford-oval/storm.git cd storm ``` 2. Install the required packages. ```shell conda create -n storm python=3.11 conda activate storm pip install -r requirements.txt ``` ## API Currently, our package support: - Language model components: All language models supported by litellm as listed [here](https://docs.litellm.ai/docs/providers) - Embedding model components: All embedding models supported by litellm as listed [here](https://docs.litellm.ai/docs/embedding/supported_embedding) - retrieval module components: `YouRM`, `BingSearch`, `VectorRM`, `SerperRM`, `BraveRM`, `SearXNG`, `DuckDuckGoSearchRM`, `TavilySearchRM`, `GoogleSearch`, and `AzureAISearch` as :star2: **PRs for integrating more search engines/retrievers into [knowledge_storm/rm.py](knowledge_storm/rm.py) are highly appreciated!** Both STORM and Co-STORM are working in the information curation layer, you need to set up the information retrieval module and language model module to create their `Runner` classes respectively. ### STORM The STORM knowledge curation engine is defined as a simple Python `STORMWikiRunner` class. Here is an example of using You.com search engine and OpenAI models. ```python 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, } # STORM is a LM system so different components can be powered by different models to reach a good balance between cost and quality. # For a good practice, choose a cheaper/faster model for `conv_simulator_lm` which is used to split queries, synthesize answers in the conversation. # Choose a more powerful model for `article_gen_lm` to generate verifiable text with citations. 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) lm_configs.set_article_polish_lm(gpt_4) # Check out the STORMWikiRunnerArguments class for more configurations. engine_args = STORMWikiRunnerArguments(...) rm = YouRM(ydc_api_key=os.getenv('YDC_API_KEY'), k=engine_args.search_top_k) runner = STORMWikiRunner(engine_args, lm_configs, rm) ``` The `STORMWikiRunner` instance can be evoked with the simple `run` method: ```python topic = input('Topic: ') runner.run( topic=topic, do_research=True, do_generate_outline=True, do_generate_article=True, do_polish_article=True, ) runner.post_run() runner.summary() ``` - `do_research`: if True, simulate conversations with difference perspectives to collect information about the topic; otherwise, load the results. - `do_generate_outline`: if True, generate an outline for the topic; otherwise, load the results. - `do_generate_article`: if True, generate an article for the topic based on the outline and the collected information; otherwise, load the results. - `do_polish_article`: if True, polish the article by adding a summarization section and (optionally) removing duplicate content; otherwise, load the results. ### Co-STORM The Co-STORM knowledge curation engine is defined as a simple Python `CoStormRunner` class. Here is an example of using Bing search engine and OpenAI models. ```python from knowledge_storm.collaborative_storm.engine import CollaborativeStormLMConfigs, RunnerArgument, CoStormRunner from knowledge_storm.lm import LitellmModel from knowledge_storm.logging_wrapper import LoggingWrapper from knowledge_storm.rm import BingSearch # Co-STORM adopts the same multi LM system paradigm as STORM lm_config: CollaborativeStormLMConfigs = CollaborativeStormLMConfigs() openai_kwargs = { "api_key": os.getenv("OPENAI_API_KEY"), "api_provider": "openai", "temperature": 1.0, "top_p": 0.9, "api_base": None, } question_answering_lm = LitellmModel(model=gpt_4o_model_name, max_tokens=1000, **openai_kwargs) discourse_manage_lm = LitellmModel(model=gpt_4o_model_name, max_tokens=500, **openai_kwargs) utterance_polishing_lm = LitellmModel(model=gpt_4o_model_name, max_tokens=2000, **openai_kwargs) warmstart_outline_gen_lm = LitellmModel(model=gpt_4o_model_name, max_tokens=500, **openai_kwargs) question_asking_lm = LitellmModel(model=gpt_4o_model_name, max_tokens=300, **openai_kwargs) knowledge_base_lm = LitellmModel(model=gpt_4o_model_name, max_tokens=1000, **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) lm_config.set_warmstart_outline_gen_lm(warmstart_outline_gen_lm) lm_config.set_question_asking_lm(question_asking_lm) lm_config.set_knowledge_base_lm(knowledge_base_lm) # Check out the Co-STORM's RunnerArguments class for more configurations. topic = input('Topic: ') runner_argument = RunnerArgument(topic=topic, ...) logging_wrapper = LoggingWrapper(lm_config) bing_rm = BingSearch(bing_search_api_key=os.environ.get("BING_SEARCH_API_KEY"), k=runner_argument.retrieve_top_k) costorm_runner = CoStormRunner(lm_config=lm_config, runner_argument=runner_argument, logging_wrapper=logging_wrapper, rm=bing_rm) ``` The `CoStormRunner` instance can be evoked with the `warmstart()` and `step(...)` methods. ```python # Warm start the system to build shared conceptual space between Co-STORM and users costorm_runner.warm_start() # Step through the collaborative discourse # Run either of the code snippets below in any order, as many times as you'd like # To observe the conversation: conv_turn = costorm_runner.step() # To inject your utterance to actively steer the conversation: costorm_runner.step(user_utterance="YOUR UTTERANCE HERE") # Generate report based on the collaborative discourse costorm_runner.knowledge_base.reorganize() article = costorm_runner.generate_report() print(article) ``` ## Quick Start with Example Scripts We provide scripts in our [examples folder](examples) as a quick start to run STORM and Co-STORM with different configurations. We suggest using `secrets.toml` to set up the API keys. Create a file `secrets.toml` under the root directory and add the following content: ```shell # ============ language model configurations ============ # Set up OpenAI API key. OPENAI_API_KEY="your_openai_api_key" # If you are using the API service provided by OpenAI, include the following line: OPENAI_API_TYPE="openai" # If you are using the API service provided by Microsoft Azure, include the following lines: OPENAI_API_TYPE="azure" AZURE_API_BASE="your_azure_api_base_url" AZURE_API_VERSION="your_azure_api_version" # ============ retriever configurations ============ BING_SEARCH_API_KEY="your_bing_search_api_key" # if using bing search # ============ encoder configurations ============ ENCODER_API_TYPE="openai" # if using openai encoder ``` ### STORM examples **To run STORM with `gpt` family models with default configurations:** Run the following command. ```bash python examples/storm_examples/run_storm_wiki_gpt.py \ --output-dir $OUTPUT_DIR \ --retriever bing \ --do-research \ --do-generate-outline \ --do-generate-article \ --do-polish-article ``` **To run STORM using your favorite language models or grounding on your own corpus:** Check out [examples/storm_examples/README.md](examples/storm_examples/README.md). ### Co-STORM examples To run Co-STORM with `gpt` family models with default configurations, 1. Add `BING_SEARCH_API_KEY="xxx"` and `ENCODER_API_TYPE="xxx"` to `secrets.toml` 2. Run the following command ```bash python examples/costorm_examples/run_costorm_gpt.py \ --output-dir $OUTPUT_DIR \ --retriever bing ``` ## Customization of the Pipeline ### STORM If you have installed the source code, you can customize STORM based on your own use case. STORM engine consists of 4 modules: 1. Knowledge Curation Module: Collects a broad coverage of information about the given topic. 2. Outline Generation Module: Organizes the collected information by generating a hierarchical outline for the curated knowledge. 3. Article Generation Module: Populates the generated outline with the collected information. 4. Article Polishing Module: Refines and enhances the written article for better presentation. The interface for each module is defined in `knowledge_storm/interface.py`, while their implementations are instantiated in `knowledge_storm/storm_wiki/modules/*`. These modules can be customized according to your specific requirements (e.g., generating sections in bullet point format instead of full paragraphs). ### Co-STORM If you have installed the source code, you can customize Co-STORM based on your own use case 1. Co-STORM introduces multiple LLM agent types (i.e. Co-STORM experts and Moderator). LLM agent interface is defined in `knowledge_storm/interface.py` , while its implementation is instantiated in `knowledge_storm/collaborative_storm/modules/co_storm_agents.py`. Different LLM agent policies can be customized. 2. Co-STORM introduces a collaborative discourse protocol, with its core function centered on turn policy management. We provide an example implementation of turn policy management through `DiscourseManager` in `knowledge_storm/collaborative_storm/engine.py`. It can be customized and further improved. ## Datasets To facilitate the study of automatic knowledge curation and complex information seeking, our project releases the following datasets: ### FreshWiki The FreshWiki Dataset is a collection of 100 high-quality Wikipedia articles focusing on the most-edited pages from February 2022 to September 2023. See Section 2.1 in [STORM paper](https://arxiv.org/abs/2402.14207) for more details. You can download the dataset from [huggingface](https://huggingface.co/datasets/EchoShao8899/FreshWiki) directly. To ease the data contamination issue, we archive the [source code](https://github.com/stanford-oval/storm/tree/NAACL-2024-code-backup/FreshWiki) for the data construction pipeline that can be repeated at future dates. ### WildSeek To study users’ interests in complex information seeking tasks in the wild, we utilized data collected from the web research preview to create the WildSeek dataset. We downsampled the data to ensure the diversity of the topics and the quality of the data. Each data point is a pair comprising a topic and the user’s goal for conducting deep search on the topic. For more details, please refer to Section 2.2 and Appendix A of [Co-STORM paper](https://www.arxiv.org/abs/2408.15232). The WildSeek dataset is available [here](https://huggingface.co/datasets/YuchengJiang/WildSeek). ## Replicate STORM & Co-STORM paper result For STORM paper experiments, please switch to the branch `NAACL-2024-code-backup` [here](https://github.com/stanford-oval/storm/tree/NAACL-2024-code-backup). For Co-STORM paper experiments, please switch to the branch `EMNLP-2024-code-backup` (placeholder for now, will be updated soon). ## Roadmap & Contributions Our team is actively working on: 1. Human-in-the-Loop Functionalities: Supporting user participation in the knowledge curation process. 2. Information Abstraction: Developing abstractions for curated information to support presentation formats beyond the Wikipedia-style report. If you have any questions or suggestions, please feel free to open an issue or pull request. We welcome contributions to improve the system and the codebase! Contact person: [Yijia Shao](mailto:shaoyj@stanford.edu) and [Yucheng Jiang](mailto:yuchengj@stanford.edu) ## Acknowledgement We would like to thank Wikipedia for its excellent open-source content. The FreshWiki dataset is sourced from Wikipedia, licensed under the Creative Commons Attribution-ShareAlike (CC BY-SA) license. We are very grateful to [Michelle Lam](https://michelle123lam.github.io/) for designing the logo for this project and [Dekun Ma](https://dekun.me) for leading the UI development. Thanks to Vercel for their support of [open-source software](https://storm.genie.stanford.edu) ## Citation Please cite our paper if you use this code or part of it in your work: ```bibtex @inproceedings{jiang-etal-2024-unknown, title = "Into the Unknown Unknowns: Engaged Human Learning through Participation in Language Model Agent Conversations", author = "Jiang, Yucheng and Shao, Yijia and Ma, Dekun and Semnani, Sina and Lam, Monica", editor = "Al-Onaizan, Yaser and Bansal, Mohit and Chen, Yun-Nung", booktitle = "Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing", month = nov, year = "2024", address = "Miami, Florida, USA", publisher = "Association for Computational Linguistics", url = "https://aclanthology.org/2024.emnlp-main.554/", doi = "10.18653/v1/2024.emnlp-main.554", pages = "9917--9955", } @inproceedings{shao-etal-2024-assisting, title = "Assisting in Writing {W}ikipedia-like Articles From Scratch with Large Language Models", author = "Shao, Yijia and Jiang, Yucheng and Kanell, Theodore and Xu, Peter and Khattab, Omar and Lam, Monica", editor = "Duh, Kevin and Gomez, Helena and Bethard, Steven", booktitle = "Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)", month = jun, year = "2024", address = "Mexico City, Mexico", publisher = "Association for Computational Linguistics", url = "https://aclanthology.org/2024.naacl-long.347/", doi = "10.18653/v1/2024.naacl-long.347", pages = "6252--6278", } ``` ================================================ FILE: examples/costorm_examples/run_costorm_gpt.py ================================================ """ Co-STORM pipeline powered by GPT-4o/4o-mini and Bing search engine. You need to set up the following environment variables to run this script: - OPENAI_API_KEY: OpenAI API key - OPENAI_API_TYPE: OpenAI API type (e.g., 'openai' or 'azure') - AZURE_API_BASE: Azure API base URL if using Azure API - AZURE_API_VERSION: Azure API version if using Azure API - BING_SEARCH_API_KEY: Biang search API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key Output will be structured as below args.output_dir/ log.json # Log of information-seeking conversation report.txt # Final article generated """ import os import json from argparse import ArgumentParser from knowledge_storm.collaborative_storm.engine import ( CollaborativeStormLMConfigs, RunnerArgument, CoStormRunner, ) from knowledge_storm.collaborative_storm.modules.callback import ( LocalConsolePrintCallBackHandler, ) from knowledge_storm.lm import OpenAIModel, AzureOpenAIModel from knowledge_storm.logging_wrapper import LoggingWrapper from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, ) from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_config: CollaborativeStormLMConfigs = CollaborativeStormLMConfigs() openai_kwargs = ( { "api_key": os.getenv("OPENAI_API_KEY"), "api_provider": "openai", "temperature": 1.0, "top_p": 0.9, "api_base": None, } if os.getenv("OPENAI_API_TYPE") == "openai" else { "api_key": os.getenv("AZURE_API_KEY"), "temperature": 1.0, "top_p": 0.9, "api_base": os.getenv("AZURE_API_BASE"), "api_version": os.getenv("AZURE_API_VERSION"), } ) ModelClass = ( OpenAIModel if os.getenv("OPENAI_API_TYPE") == "openai" else AzureOpenAIModel ) # If you are using Azure service, make sure the model name matches your own deployed model name. # The default name here is only used for demonstration and may not match your case. gpt_4o_mini_model_name = "gpt-4o-mini" gpt_4o_model_name = "gpt-4o" if os.getenv("OPENAI_API_TYPE") == "azure": openai_kwargs["api_base"] = os.getenv("AZURE_API_BASE") openai_kwargs["api_version"] = os.getenv("AZURE_API_VERSION") # STORM is a LM system so different components can be powered by different models. # For a good balance between cost and quality, you can choose a cheaper/faster model for conv_simulator_lm # which is used to split queries, synthesize answers in the conversation. We recommend using stronger models # for outline_gen_lm which is responsible for organizing the collected information, and article_gen_lm # which is responsible for generating sections with citations. question_answering_lm = ModelClass( model=gpt_4o_model_name, max_tokens=1000, **openai_kwargs ) discourse_manage_lm = ModelClass( model=gpt_4o_model_name, max_tokens=500, **openai_kwargs ) utterance_polishing_lm = ModelClass( model=gpt_4o_model_name, max_tokens=2000, **openai_kwargs ) warmstart_outline_gen_lm = ModelClass( model=gpt_4o_model_name, max_tokens=500, **openai_kwargs ) question_asking_lm = ModelClass( model=gpt_4o_model_name, max_tokens=300, **openai_kwargs ) knowledge_base_lm = ModelClass( model=gpt_4o_model_name, max_tokens=1000, **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) lm_config.set_warmstart_outline_gen_lm(warmstart_outline_gen_lm) lm_config.set_question_asking_lm(question_asking_lm) lm_config.set_knowledge_base_lm(knowledge_base_lm) topic = input("Topic: ") runner_argument = RunnerArgument( topic=topic, retrieve_top_k=args.retrieve_top_k, max_search_queries=args.max_search_queries, total_conv_turn=args.total_conv_turn, max_search_thread=args.max_search_thread, max_search_queries_per_turn=args.max_search_queries_per_turn, warmstart_max_num_experts=args.warmstart_max_num_experts, warmstart_max_turn_per_experts=args.warmstart_max_turn_per_experts, warmstart_max_thread=args.warmstart_max_thread, max_thread_num=args.max_thread_num, max_num_round_table_experts=args.max_num_round_table_experts, moderator_override_N_consecutive_answering_turn=args.moderator_override_N_consecutive_answering_turn, node_expansion_trigger_count=args.node_expansion_trigger_count, ) logging_wrapper = LoggingWrapper(lm_config) callback_handler = ( LocalConsolePrintCallBackHandler() if args.enable_log_print else None ) # Co-STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=runner_argument.retrieve_top_k, ) case "you": rm = YouRM( ydc_api_key=os.getenv("YDC_API_KEY"), k=runner_argument.retrieve_top_k ) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=runner_argument.retrieve_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=runner_argument.retrieve_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=runner_argument.retrieve_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=runner_argument.retrieve_top_k, ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", or "searxng"' ) costorm_runner = CoStormRunner( lm_config=lm_config, runner_argument=runner_argument, logging_wrapper=logging_wrapper, rm=rm, callback_handler=callback_handler, ) # warm start the system costorm_runner.warm_start() # Below is an example of how users may interact with Co-STORM to seek information together # In actual deployment, we suggest allowing the user to decide whether to observe the agent utterance or inject a turn # observing Co-STORM LLM agent utterance for 5 turns for _ in range(1): conv_turn = costorm_runner.step() print(f"**{conv_turn.role}**: {conv_turn.utterance}\n") # active engaging by injecting your utterance your_utterance = input("Your utterance: ") costorm_runner.step(user_utterance=your_utterance) # continue observing conv_turn = costorm_runner.step() print(f"**{conv_turn.role}**: {conv_turn.utterance}\n") # generate report costorm_runner.knowledge_base.reorganize() article = costorm_runner.generate_report() # save results os.makedirs(args.output_dir, exist_ok=True) # Save article with open(os.path.join(args.output_dir, "report.md"), "w") as f: f.write(article) # Save instance dump instance_copy = costorm_runner.to_dict() with open(os.path.join(args.output_dir, "instance_dump.json"), "w") as f: json.dump(instance_copy, f, indent=2) # Save logging log_dump = costorm_runner.dump_logging_and_reset() with open(os.path.join(args.output_dir, "log.json"), "w") as f: json.dump(log_dump, f, indent=2) if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/co-storm", help="Directory to store the outputs.", ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng"], help="The search engine API to use for retrieving information.", ) # hyperparameters for co-storm parser.add_argument( "--retrieve_top_k", type=int, default=10, help="Retrieve top k results for each query in retriever.", ) parser.add_argument( "--max_search_queries", type=int, default=2, help="Maximum number of search queries to consider for each question.", ) parser.add_argument( "--total_conv_turn", type=int, default=20, help="Maximum number of turns in conversation.", ) parser.add_argument( "--max_search_thread", type=int, default=5, help="Maximum number of parallel threads for retriever.", ) parser.add_argument( "--max_search_queries_per_turn", type=int, default=3, help="Maximum number of search queries to consider in each turn.", ) parser.add_argument( "--warmstart_max_num_experts", type=int, default=3, help="Max number of experts in perspective-guided QA during warm start.", ) parser.add_argument( "--warmstart_max_turn_per_experts", type=int, default=2, help="Max number of turns per perspective during warm start.", ) parser.add_argument( "--warmstart_max_thread", type=int, default=3, help="Max number of threads for parallel perspective-guided QA during warm start.", ) parser.add_argument( "--max_thread_num", type=int, default=10, help=( "Maximum number of threads to use. " "Consider reducing it if you keep getting 'Exceed rate limit' errors when calling the LM API." ), ) parser.add_argument( "--max_num_round_table_experts", type=int, default=2, help="Max number of active experts in round table discussion.", ) parser.add_argument( "--moderator_override_N_consecutive_answering_turn", type=int, default=3, help=( "Number of consecutive expert answering turns before the moderator overrides the conversation." ), ) parser.add_argument( "--node_expansion_trigger_count", type=int, default=10, help="Trigger node expansion for nodes that contain more than N snippets.", ) # Boolean flags parser.add_argument( "--enable_log_print", action="store_true", help="If set, enable console log print.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/README.md ================================================ # Examples We host a number of example scripts for various customization of STORM (e.g., use your favorite language models, use your own corpus, etc.). These examples can be starting points for your own customizations and you are welcome to contribute your own examples by submitting a pull request to this directory. ## Run STORM with your own language model [run_storm_wiki_gpt.py](run_storm_wiki_gpt.py) provides an example of running STORM with GPT models, and [run_storm_wiki_claude.py](run_storm_wiki_claude.py) provides an example of running STORM with Claude models. Besides using close-source models, you can also run STORM with models with open weights. `run_storm_wiki_mistral.py` provides an example of running STORM with `Mistral-7B-Instruct-v0.2` using [VLLM](https://docs.vllm.ai/en/stable/) server: 1. Set up a VLLM server with the `Mistral-7B-Instruct-v0.2` model running. 2. Run the following command under the root directory of the repository: ``` python examples/storm_examples/run_storm_wiki_mistral.py \ --url $URL \ --port $PORT \ --output-dir $OUTPUT_DIR \ --retriever you \ --do-research \ --do-generate-outline \ --do-generate-article \ --do-polish-article ``` - `--url` URL of the VLLM server. - `--port` Port of the VLLM server. Besides VLLM server, STORM is also compatible with [TGI](https://huggingface.co/docs/text-generation-inference/en/index) server or [Together.ai](https://www.together.ai/products#inference) endpoint. ## Run STORM with your own corpus By default, STORM is grounded on the Internet using the search engine, but it can also be grounded on your own corpus using `VectorRM`. [run_storm_wiki_with_gpt_with_VectorRM.py](run_storm_wiki_gpt_with_VectorRM.py) provides an example of running STORM grounding on your provided data. 1. Set up API keys. - Make sure you have set up the OpenAI API key. - `VectorRM` use [Qdrant](https://github.com/qdrant/qdrant-client) to create a vector store. If you want to set up this vector store online on a [Qdrant cloud server](https://cloud.qdrant.io/login), you need to set up `QDRANT_API_KEY` in `secrets.toml` as well; if you want to save the vector store locally, make sure you provide a location for the vector store. 2. Prepare your corpus. The documents should be provided as a single CSV file with the following format: | content | title | url | description | |------------------------|------------|------------|------------------------------------| | I am a document. | Document 1 | docu-n-112 | A self-explanatory document. | | I am another document. | Document 2 | docu-l-13 | Another self-explanatory document. | | ... | ... | ... | ... | - `url` will be used as a unique identifier of the document in STORM engine, so ensure different documents have different urls. - The contents for `title` and `description` columns are optional. If not provided, the script will use default empty values. - The content column is crucial and should be provided for each document. 3. Run the command under the root directory of the repository: To create the vector store offline, run ``` python examples/storm_examples/run_storm_wiki_gpt_with_VectorRM.py \ --output-dir $OUTPUT_DIR \ --vector-db-mode offline \ --offline-vector-db-dir $OFFLINE_VECTOR_DB_DIR \ --csv-file-path $CSV_FILE_PATH \ --device $DEVICE_FOR_EMBEDDING(mps, cuda, cpu) \ --do-research \ --do-generate-outline \ --do-generate-article \ --do-polish-article ``` To create the vector store online on a Qdrant server, run ``` python examples/storm_examples/run_storm_wiki_gpt_with_VectorRM.py \ --output-dir $OUTPUT_DIR \ --vector-db-mode online \ --online-vector-db-url $ONLINE_VECTOR_DB_URL \ --csv-file-path $CSV_FILE_PATH \ --device $DEVICE_FOR_EMBEDDING(mps, cuda, cpu) \ --do-research \ --do-generate-outline \ --do-generate-article \ --do-polish-article ``` 4. **Quick test with Kaggle arXiv Paper Abstracts dataset**: - Download `arxiv_data_210930-054931.csv` from [here](https://www.kaggle.com/datasets/spsayakpaul/arxiv-paper-abstracts). - Run the following command under the root directory to downsample the dataset by filtering papers with terms `[cs.CV]` and get a csv file that match the format mentioned above. ``` python examples/storm_examples/helper/process_kaggle_arxiv_abstract_dataset.py --input-path $PATH_TO_THE_DOWNLOADED_FILE --output-path $PATH_TO_THE_PROCESSED_CSV ``` - Run the following command to run STORM grounding on the processed dataset. You can input a topic related to computer vision (e.g., "The progress of multimodal models in computer vision") to see the generated article. (Note that the generated article may not include enough details since the quick test only use the abstracts of arxiv papers.) ``` python examples/storm_examples/run_storm_wiki_gpt_with_VectorRM.py \ --output-dir $OUTPUT_DIR \ --vector-db-mode offline \ --offline-vector-db-dir $OFFLINE_VECTOR_DB_DIR \ --csv-file-path $PATH_TO_THE_PROCESSED_CSV \ --device $DEVICE_FOR_EMBEDDING(mps, cuda, cpu) \ --do-research \ --do-generate-outline \ --do-generate-article \ --do-polish-article ``` - For a quicker run, you can also download the pre-embedded vector store directly from [here](https://drive.google.com/file/d/1bijFkw5BKU7bqcmXMhO-5hg2fdKAL9bf/view?usp=share_link). ``` python examples/storm_examples/run_storm_wiki_gpt_with_VectorRM.py \ --output-dir $OUTPUT_DIR \ --vector-db-mode offline \ --offline-vector-db-dir $DOWNLOADED_VECTOR_DB_DR \ --do-research \ --do-generate-outline \ --do-generate-article \ --do-polish-article ``` ================================================ FILE: examples/storm_examples/helper/process_kaggle_arxiv_abstract_dataset.py ================================================ """Process `arxiv_data_210930-054931.csv` from https://www.kaggle.com/datasets/spsayakpaul/arxiv-paper-abstracts to a csv file that is compatible with VectorRM. """ from argparse import ArgumentParser import pandas as pd if __name__ == "__main__": parser = ArgumentParser() parser.add_argument( "--input-path", type=str, help="Path to arxiv_data_210930-054931.csv." ) parser.add_argument( "--output-path", type=str, help="Path to store the csv file that is compatible with VectorRM.", ) args = parser.parse_args() df = pd.read_csv(args.input_path) print(f"The original dataset has {len(df)} samples.") # Downsample the dataset. df = df[df["terms"] == "['cs.CV']"] # Reformat the dataset to match the VectorRM input format. df.rename(columns={"abstracts": "content", "titles": "title"}, inplace=True) df["url"] = [ "uid_" + str(idx) for idx in range(len(df)) ] # Ensure the url is unique. df["description"] = "" print(f"The downsampled dataset has {len(df)} samples.") df.to_csv(args.output_path, index=False) ================================================ FILE: examples/storm_examples/run_storm_wiki_claude.py ================================================ """ STORM Wiki pipeline powered by Claude family models and You.com search engine. You need to set up the following environment variables to run this script: - ANTHROPIC_API_KEY: Anthropic API key - YDC_API_KEY: You.com API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os from argparse import ArgumentParser from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import ClaudeModel from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, ) from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() claude_kwargs = { "api_key": os.getenv("ANTHROPIC_API_KEY"), "temperature": 1.0, "top_p": 0.9, } # STORM is a LM system so different components can be powered by different models. # For a good balance between cost and quality, you can choose a cheaper/faster model for conv_simulator_lm # which is used to split queries, synthesize answers in the conversation. We recommend using stronger models # for outline_gen_lm which is responsible for organizing the collected information, and article_gen_lm # which is responsible for generating sections with citations. conv_simulator_lm = ClaudeModel( model="claude-3-haiku-20240307", max_tokens=500, **claude_kwargs ) question_asker_lm = ClaudeModel( model="claude-3-sonnet-20240229", max_tokens=500, **claude_kwargs ) outline_gen_lm = ClaudeModel( model="claude-3-opus-20240229", max_tokens=400, **claude_kwargs ) article_gen_lm = ClaudeModel( model="claude-3-opus-20240229", max_tokens=700, **claude_kwargs ) article_polish_lm = ClaudeModel( model="claude-3-opus-20240229", max_tokens=4000, **claude_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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case "you": rm = YouRM(ydc_api_key=os.getenv("YDC_API_KEY"), k=engine_args.search_top_k) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=engine_args.search_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=engine_args.search_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=engine_args.search_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", or "searxng"' ) runner = STORMWikiRunner(engine_args, lm_configs, rm) topic = input("Topic: ") runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/claude", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng"], help="The search engine API to use for retrieving information.", ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_deepseek.py ================================================ """ STORM Wiki pipeline powered by DeepSeek models and You.com or Bing search engine. You need to set up the following environment variables to run this script: - DEEPSEEK_API_KEY: DeepSeek API key - DEEPSEEK_API_BASE: DeepSeek API base URL (default is https://api.deepseek.com) - YDC_API_KEY: You.com API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os import re import logging from argparse import ArgumentParser from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import DeepSeekModel from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, ) from knowledge_storm.utils import load_api_key def sanitize_topic(topic): """ Sanitize the topic name for use in file names. Remove or replace characters that are not allowed in file names. """ # Replace spaces with underscores topic = topic.replace(" ", "_") # Remove any character that isn't alphanumeric, underscore, or hyphen topic = re.sub(r"[^a-zA-Z0-9_-]", "", topic) # Ensure the topic isn't empty after sanitization if not topic: topic = "unnamed_topic" return topic def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() logger = logging.getLogger(__name__) # Ensure DEEPSEEK_API_KEY is set if not os.getenv("DEEPSEEK_API_KEY"): raise ValueError( "DEEPSEEK_API_KEY environment variable is not set. Please set it in your secrets.toml file." ) deepseek_kwargs = { "api_key": os.getenv("DEEPSEEK_API_KEY"), "api_base": os.getenv("DEEPSEEK_API_BASE", "https://api.deepseek.com"), "temperature": args.temperature, "top_p": args.top_p, } # 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=args.model, max_tokens=500, **deepseek_kwargs ) question_asker_lm = DeepSeekModel( model=args.model, max_tokens=500, **deepseek_kwargs ) outline_gen_lm = DeepSeekModel(model=args.model, max_tokens=400, **deepseek_kwargs) article_gen_lm = DeepSeekModel(model=args.model, max_tokens=700, **deepseek_kwargs) article_polish_lm = DeepSeekModel( model=args.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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case "you": rm = YouRM(ydc_api_key=os.getenv("YDC_API_KEY"), k=engine_args.search_top_k) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=engine_args.search_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=engine_args.search_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=engine_args.search_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", or "searxng"' ) runner = STORMWikiRunner(engine_args, lm_configs, rm) topic = input("Topic: ") sanitized_topic = sanitize_topic(topic) try: runner.run( topic=sanitized_topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, remove_duplicate=args.remove_duplicate, ) runner.post_run() runner.summary() except Exception as e: logger.exception(f"An error occurred: {str(e)}") raise if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/deepseek", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng"], help="The search engine API to use for retrieving information.", ) parser.add_argument( "--model", type=str, choices=["deepseek-chat", "deepseek-coder"], default="deepseek-chat", help='DeepSeek model to use. "deepseek-chat" for general tasks, "deepseek-coder" for coding tasks.', ) parser.add_argument( "--temperature", type=float, default=1.0, help="Sampling temperature to use." ) parser.add_argument( "--top_p", type=float, default=0.9, help="Top-p sampling parameter." ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_gemini.py ================================================ """ STORM Wiki pipeline powered by Google Gemini models and search engine. You need to set up the following environment variables to run this script: - GOOGLE_API_KEY: Google API key (Can be obtained from https://ai.google.dev/gemini-api/docs/api-key) - YDC_API_KEY: You.com API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os from argparse import ArgumentParser from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import GoogleModel from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, ) from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() gemini_kwargs = { "api_key": os.getenv("GOOGLE_API_KEY"), "temperature": 1.0, "top_p": 0.9, } # STORM is a LM system so different components can be powered by different models. # For a good balance between cost and quality, you can choose a cheaper/faster model for conv_simulator_lm # which is used to split queries, synthesize answers in the conversation. We recommend using stronger models # for outline_gen_lm which is responsible for organizing the collected information, and article_gen_lm # which is responsible for generating sections with citations. # To check out available Google models, see: # https://ai.google.dev/gemini-api/docs/get-started/tutorial?lang=python#list_models conv_simulator_lm = GoogleModel( model="models/gemini-1.5-flash", max_tokens=500, **gemini_kwargs ) question_asker_lm = GoogleModel( model="models/gemini-1.5-flash", max_tokens=500, **gemini_kwargs ) outline_gen_lm = GoogleModel( model="models/gemini-1.5-pro-exp-0801", max_tokens=400, **gemini_kwargs ) article_gen_lm = GoogleModel( model="models/gemini-1.5-pro-exp-0801", max_tokens=700, **gemini_kwargs ) article_polish_lm = GoogleModel( model="models/gemini-1.5-pro-exp-0801", max_tokens=4000, **gemini_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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case "you": rm = YouRM(ydc_api_key=os.getenv("YDC_API_KEY"), k=engine_args.search_top_k) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=engine_args.search_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=engine_args.search_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=engine_args.search_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", or "searxng"' ) runner = STORMWikiRunner(engine_args, lm_configs, rm) topic = input("Topic: ") runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/gemini", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng"], help="The search engine API to use for retrieving information.", ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_gpt.py ================================================ """ STORM Wiki pipeline powered by GPT-3.5/4 and You.com search engine. You need to set up the following environment variables to run this script: - OPENAI_API_KEY: OpenAI API key - OPENAI_API_TYPE: OpenAI API type (e.g., 'openai' or 'azure') - AZURE_API_BASE: Azure API base URL if using Azure API - AZURE_API_VERSION: Azure API version if using Azure API - YDC_API_KEY: You.com API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os from argparse import ArgumentParser from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import OpenAIModel, AzureOpenAIModel from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, AzureAISearch, ) from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() openai_kwargs = { "api_key": os.getenv("OPENAI_API_KEY"), "temperature": 1.0, "top_p": 0.9, } ModelClass = ( OpenAIModel if os.getenv("OPENAI_API_TYPE") == "openai" else AzureOpenAIModel ) # If you are using Azure service, make sure the model name matches your own deployed model name. # The default name here is only used for demonstration and may not match your case. gpt_35_model_name = ( "gpt-3.5-turbo" if os.getenv("OPENAI_API_TYPE") == "openai" else "gpt-35-turbo" ) gpt_4_model_name = "gpt-4o" if os.getenv("OPENAI_API_TYPE") == "azure": openai_kwargs["api_base"] = os.getenv("AZURE_API_BASE") openai_kwargs["api_version"] = os.getenv("AZURE_API_VERSION") # STORM is a LM system so different components can be powered by different models. # For a good balance between cost and quality, you can choose a cheaper/faster model for conv_simulator_lm # which is used to split queries, synthesize answers in the conversation. We recommend using stronger models # for outline_gen_lm which is responsible for organizing the collected information, and article_gen_lm # which is responsible for generating sections with citations. 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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case "you": rm = YouRM(ydc_api_key=os.getenv("YDC_API_KEY"), k=engine_args.search_top_k) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=engine_args.search_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=engine_args.search_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=engine_args.search_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k ) case "azure_ai_search": rm = AzureAISearch( azure_ai_search_api_key=os.getenv("AZURE_AI_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", "searxng", or "azure_ai_search"' ) runner = STORMWikiRunner(engine_args, lm_configs, rm) topic = input("Topic: ") runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/gpt", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=[ "bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng", "azure_ai_search", ], help="The search engine API to use for retrieving information.", ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_gpt_with_VectorRM.py ================================================ """ This STORM Wiki pipeline powered by GPT-3.5/4 and local retrieval model that uses Qdrant. You need to set up the following environment variables to run this script: - OPENAI_API_KEY: OpenAI API key - OPENAI_API_TYPE: OpenAI API type (e.g., 'openai' or 'azure') - QDRANT_API_KEY: Qdrant API key (needed ONLY if online vector store was used) You will also need an existing Qdrant vector store either saved in a folder locally offline or in a server online. If not, then you would need a CSV file with documents, and the script is going to create the vector store for you. The CSV should be in the following format: content | title | url | description I am a document. | Document 1 | docu-n-112 | A self-explanatory document. I am another document. | Document 2 | docu-l-13 | Another self-explanatory document. Notice that the URL will be a unique identifier for the document so ensure different documents have different urls. Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os from argparse import ArgumentParser from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.rm import VectorRM from knowledge_storm.lm import OpenAIModel, AzureOpenAIModel from knowledge_storm.utils import load_api_key, QdrantVectorStoreManager def main(args): # Load API key from the specified toml file path load_api_key(toml_file_path="secrets.toml") # Initialize the language model configurations engine_lm_configs = STORMWikiLMConfigs() openai_kwargs = { "api_key": os.getenv("OPENAI_API_KEY"), "temperature": 1.0, "top_p": 0.9, } ModelClass = ( OpenAIModel if os.getenv("OPENAI_API_TYPE") == "openai" else AzureOpenAIModel ) # If you are using Azure service, make sure the model name matches your own deployed model name. # The default name here is only used for demonstration and may not match your case. gpt_35_model_name = ( "gpt-3.5-turbo" if os.getenv("OPENAI_API_TYPE") == "openai" else "gpt-35-turbo" ) gpt_4_model_name = "gpt-4o" if os.getenv("OPENAI_API_TYPE") == "azure": openai_kwargs["api_base"] = os.getenv("AZURE_API_BASE") openai_kwargs["api_version"] = os.getenv("AZURE_API_VERSION") # STORM is a LM system so different components can be powered by different models. # For a good balance between cost and quality, you can choose a cheaper/faster model for conv_simulator_lm # which is used to split queries, synthesize answers in the conversation. We recommend using stronger models # for outline_gen_lm which is responsible for organizing the collected information, and article_gen_lm # which is responsible for generating sections with citations. 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 ) engine_lm_configs.set_conv_simulator_lm(conv_simulator_lm) engine_lm_configs.set_question_asker_lm(question_asker_lm) engine_lm_configs.set_outline_gen_lm(outline_gen_lm) engine_lm_configs.set_article_gen_lm(article_gen_lm) engine_lm_configs.set_article_polish_lm(article_polish_lm) # Initialize the engine arguments engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # Create / update the vector store with the documents in the csv file if args.csv_file_path: kwargs = { "file_path": args.csv_file_path, "content_column": "content", "title_column": "title", "url_column": "url", "desc_column": "description", "batch_size": args.embed_batch_size, "vector_db_mode": args.vector_db_mode, "collection_name": args.collection_name, "embedding_model": args.embedding_model, "device": args.device, } if args.vector_db_mode == "offline": QdrantVectorStoreManager.create_or_update_vector_store( vector_store_path=args.offline_vector_db_dir, **kwargs ) elif args.vector_db_mode == "online": QdrantVectorStoreManager.create_or_update_vector_store( url=args.online_vector_db_url, api_key=os.getenv("QDRANT_API_KEY"), **kwargs ) # Setup VectorRM to retrieve information from your own data rm = VectorRM( collection_name=args.collection_name, embedding_model=args.embedding_model, device=args.device, k=engine_args.search_top_k, ) # initialize the vector store, either online (store the db on Qdrant server) or offline (store the db locally): if args.vector_db_mode == "offline": rm.init_offline_vector_db(vector_store_path=args.offline_vector_db_dir) elif args.vector_db_mode == "online": rm.init_online_vector_db( url=args.online_vector_db_url, api_key=os.getenv("QDRANT_API_KEY") ) # Initialize the STORM Wiki Runner runner = STORMWikiRunner(engine_args, engine_lm_configs, rm) # run the pipeline topic = input("Topic: ") runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/gpt_retrieval", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) # provide local corpus and set up vector db parser.add_argument( "--collection-name", type=str, default="my_documents", help="The collection name for vector store.", ) parser.add_argument( "--embedding_model", type=str, default="BAAI/bge-m3", help="The collection name for vector store.", ) parser.add_argument( "--device", type=str, default="mps", help="The device used to run the retrieval model (mps, cuda, cpu, etc).", ) parser.add_argument( "--vector-db-mode", type=str, choices=["offline", "online"], help="The mode of the Qdrant vector store (offline or online).", ) parser.add_argument( "--offline-vector-db-dir", type=str, default="./vector_store", help="If use offline mode, please provide the directory to store the vector store.", ) parser.add_argument( "--online-vector-db-url", type=str, help="If use online mode, please provide the url of the Qdrant server.", ) parser.add_argument( "--csv-file-path", type=str, default=None, help="The path of the custom document corpus in CSV format. The CSV file should include " "content, title, url, and description columns.", ) parser.add_argument( "--embed-batch-size", type=int, default=64, help="Batch size for embedding the documents in the csv file.", ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_groq.py ================================================ """ STORM Wiki pipeline powered by llama3-70b-8192 hosted by Groq server and You.com search engine. You need to set up the following environment variables to run this script: - GROQ_API_KEY: You can get your Groq API Key at https://console.groq.com/keys - YDC_API_KEY: You.com API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key You also need to have a VLLM server running with the Mistral-7B-Instruct-v0.2 model. Specify `--url` and `--port` accordingly. Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os import re from argparse import ArgumentParser from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) # Now import lm directly import lm from lm import GroqModel from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, ) from knowledge_storm.utils import load_api_key def sanitize_topic(topic): """ Sanitize the topic name for use in file names. Remove or replace characters that are not allowed in file names. """ # Replace spaces with underscores topic = topic.replace(" ", "_") # Remove any character that isn't alphanumeric, underscore, or hyphen topic = re.sub(r"[^a-zA-Z0-9_-]", "", topic) # Ensure the topic isn't empty after sanitization if not topic: topic = "unnamed_topic" return topic def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() # Ensure GROQ_API_KEY is set if not os.getenv("GROQ_API_KEY"): raise ValueError( "GROQ_API_KEY environment variable is not set. Please set it in your secrets.toml file." ) groq_kwargs = { "api_key": os.getenv("GROQ_API_KEY"), "api_base": "https://api.groq.com/openai/v1", "temperature": args.temperature, "top_p": args.top_p, } # Groq currently offers the "llama3-70b-8192" model with generous free API credits and the llama3.1 family of models as a preview for paying customers conv_simulator_lm = GroqModel( model="llama3-70b-8192", max_tokens=500, **groq_kwargs ) question_asker_lm = GroqModel( model="llama3-70b-8192", max_tokens=500, **groq_kwargs ) outline_gen_lm = GroqModel(model="llama3-70b-8192", max_tokens=400, **groq_kwargs) article_gen_lm = GroqModel(model="llama3-70b-8192", max_tokens=700, **groq_kwargs) article_polish_lm = GroqModel( model="llama3-70b-8192", max_tokens=4000, **groq_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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case "you": rm = YouRM(ydc_api_key=os.getenv("YDC_API_KEY"), k=engine_args.search_top_k) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=engine_args.search_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=engine_args.search_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=engine_args.search_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", or "searxng"' ) runner = STORMWikiRunner(engine_args, lm_configs, rm) topic = input("Topic: ") sanitized_topic = sanitize_topic(topic) try: runner.run( topic=sanitized_topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, remove_duplicate=args.remove_duplicate, ) runner.post_run() runner.summary() except Exception as e: logger.exception(f"An error occurred: {str(e)}") raise if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/groq", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng"], help="The search engine API to use for retrieving information.", ) parser.add_argument( "--temperature", type=float, default=1.0, help="Sampling temperature to use." ) parser.add_argument( "--top_p", type=float, default=0.9, help="Top-p sampling parameter." ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_mistral.py ================================================ """ STORM Wiki pipeline powered by Mistral-7B-Instruct-v0.2 hosted by VLLM server and You.com search engine. You need to set up the following environment variables to run this script: - YDC_API_KEY: You.com API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key You also need to have a VLLM server running with the Mistral-7B-Instruct-v0.2 model. Specify `--url` and `--port` accordingly. Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os from argparse import ArgumentParser from dspy import Example from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import VLLMClient from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, ) from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() mistral_kwargs = { "model": "mistralai/Mistral-7B-Instruct-v0.2", "port": args.port, "url": args.url, "stop": ( "\n\n---", ), # dspy uses "\n\n---" to separate examples. Open models sometimes generate this. } conv_simulator_lm = VLLMClient(max_tokens=500, **mistral_kwargs) question_asker_lm = VLLMClient(max_tokens=500, **mistral_kwargs) outline_gen_lm = VLLMClient(max_tokens=400, **mistral_kwargs) article_gen_lm = VLLMClient(max_tokens=700, **mistral_kwargs) article_polish_lm = VLLMClient(max_tokens=4000, **mistral_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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case "you": rm = YouRM(ydc_api_key=os.getenv("YDC_API_KEY"), k=engine_args.search_top_k) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=engine_args.search_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=engine_args.search_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=engine_args.search_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", or "searxng"' ) runner = STORMWikiRunner(engine_args, lm_configs, rm) # Open LMs are generally weaker in following output format. # One way for mitigation is to add one-shot example to the prompt to exemplify the desired output format. # For example, we can add the following examples to the two prompts used in StormPersonaGenerator. # Note that the example should be an object of dspy.Example with fields matching the InputField # and OutputField in the prompt (i.e., dspy.Signature). find_related_topic_example = Example( topic="Knowledge Curation", related_topics="https://en.wikipedia.org/wiki/Knowledge_management\n" "https://en.wikipedia.org/wiki/Information_science\n" "https://en.wikipedia.org/wiki/Library_science\n", ) gen_persona_example = Example( topic="Knowledge Curation", examples="Title: Knowledge management\n" "Table of Contents: History\nResearch\n Dimensions\n Strategies\n Motivations\nKM technologies" "\nKnowledge barriers\nKnowledge retention\nKnowledge audit\nKnowledge protection\n" " Knowledge protection methods\n Formal methods\n Informal methods\n" " Balancing knowledge protection and knowledge sharing\n Knowledge protection risks", personas="1. Historian of Knowledge Systems: This editor will focus on the history and evolution of knowledge curation. They will provide context on how knowledge curation has changed over time and its impact on modern practices.\n" "2. Information Science Professional: With insights from 'Information science', this editor will explore the foundational theories, definitions, and philosophy that underpin knowledge curation\n" "3. Digital Librarian: This editor will delve into the specifics of how digital libraries operate, including software, metadata, digital preservation.\n" "4. Technical expert: This editor will focus on the technical aspects of knowledge curation, such as common features of content management systems.\n" "5. Museum Curator: The museum curator will contribute expertise on the curation of physical items and the transition of these practices into the digital realm.", ) runner.storm_knowledge_curation_module.persona_generator.create_writer_with_persona.find_related_topic.demos = [ find_related_topic_example ] runner.storm_knowledge_curation_module.persona_generator.create_writer_with_persona.gen_persona.demos = [ gen_persona_example ] # A trade-off of adding one-shot example is that it will increase the input length of the prompt. Also, some # examples may be very long (e.g., an example for writing a section based on the given information), which may # confuse the model. For these cases, you can create a pseudo-example that is short and easy to understand to steer # the model's output format. # For example, we can add the following pseudo-examples to the prompt used in WritePageOutlineFromConv and # ConvToSection. write_page_outline_example = Example( topic="Example Topic", conv="Wikipedia Writer: ...\nExpert: ...\nWikipedia Writer: ...\nExpert: ...", old_outline="# Section 1\n## Subsection 1\n## Subsection 2\n" "# Section 2\n## Subsection 1\n## Subsection 2\n" "# Section 3", outline="# New Section 1\n## New Subsection 1\n## New Subsection 2\n" "# New Section 2\n" "# New Section 3\n## New Subsection 1\n## New Subsection 2\n## New Subsection 3", ) runner.storm_outline_generation_module.write_outline.write_page_outline.demos = [ write_page_outline_example ] write_section_example = Example( info="[1]\nInformation in document 1\n[2]\nInformation in document 2\n[3]\nInformation in document 3", topic="Example Topic", section="Example Section", output="# Example Topic\n## Subsection 1\n" "This is an example sentence [1]. This is another example sentence [2][3].\n" "## Subsection 2\nThis is one more example sentence [1].", ) runner.storm_article_generation.section_gen.write_section.demos = [ write_section_example ] topic = input("Topic: ") runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--url", type=str, default="http://localhost", help="URL of the VLLM server." ) parser.add_argument( "--port", type=int, default=8000, help="Port of the VLLM server." ) parser.add_argument( "--output-dir", type=str, default="./results/mistral_7b", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng"], help="The search engine API to use for retrieving information.", ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_ollama.py ================================================ """ STORM Wiki pipeline powered by local model hosted by Ollama server and You.com or Bing search engine. You need to set up the following environment variables to run this script: - YDC_API_KEY: You.com API key; BING_SEARCH_API_KEY: Bing Search API key, SERPER_API_KEY: Serper API key, BRAVE_API_KEY: Brave API key, or TAVILY_API_KEY: Tavily API key You also need to have a Ollama server running with the llama3 model or other. Specify `--url`, `--port` and `--model` accordingly. Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os import sys from argparse import ArgumentParser from dspy import Example from knowledge_storm.lm import OllamaClient from knowledge_storm.rm import ( YouRM, BingSearch, BraveRM, SerperRM, DuckDuckGoSearchRM, TavilySearchRM, SearXNG, ) from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() ollama_kwargs = { "model": args.model, "port": args.port, "url": args.url, "stop": ( "\n\n---", ), # dspy uses "\n\n---" to separate examples. Open models sometimes generate this. } conv_simulator_lm = OllamaClient(max_tokens=500, **ollama_kwargs) question_asker_lm = OllamaClient(max_tokens=500, **ollama_kwargs) outline_gen_lm = OllamaClient(max_tokens=400, **ollama_kwargs) article_gen_lm = OllamaClient(max_tokens=700, **ollama_kwargs) article_polish_lm = OllamaClient(max_tokens=4000, **ollama_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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # STORM is a knowledge curation system which consumes information from the retrieval module. # Currently, the information source is the Internet and we use search engine API as the retrieval module. match args.retriever: case "bing": rm = BingSearch( bing_search_api=os.getenv("BING_SEARCH_API_KEY"), k=engine_args.search_top_k, ) case "you": rm = YouRM(ydc_api_key=os.getenv("YDC_API_KEY"), k=engine_args.search_top_k) case "brave": rm = BraveRM( brave_search_api_key=os.getenv("BRAVE_API_KEY"), k=engine_args.search_top_k, ) case "duckduckgo": rm = DuckDuckGoSearchRM( k=engine_args.search_top_k, safe_search="On", region="us-en" ) case "serper": rm = SerperRM( serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params={"autocorrect": True, "num": 10, "page": 1}, ) case "tavily": rm = TavilySearchRM( tavily_search_api_key=os.getenv("TAVILY_API_KEY"), k=engine_args.search_top_k, include_raw_content=True, ) case "searxng": rm = SearXNG( searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k ) case _: raise ValueError( f'Invalid retriever: {args.retriever}. Choose either "bing", "you", "brave", "duckduckgo", "serper", "tavily", or "searxng"' ) runner = STORMWikiRunner(engine_args, lm_configs, rm) # Open LMs are generally weaker in following output format. # One way for mitigation is to add one-shot example to the prompt to exemplify the desired output format. # For example, we can add the following examples to the two prompts used in StormPersonaGenerator. # Note that the example should be an object of dspy.Example with fields matching the InputField # and OutputField in the prompt (i.e., dspy.Signature). find_related_topic_example = Example( topic="Knowledge Curation", related_topics="https://en.wikipedia.org/wiki/Knowledge_management\n" "https://en.wikipedia.org/wiki/Information_science\n" "https://en.wikipedia.org/wiki/Library_science\n", ) gen_persona_example = Example( topic="Knowledge Curation", examples="Title: Knowledge management\n" "Table of Contents: History\nResearch\n Dimensions\n Strategies\n Motivations\nKM technologies" "\nKnowledge barriers\nKnowledge retention\nKnowledge audit\nKnowledge protection\n" " Knowledge protection methods\n Formal methods\n Informal methods\n" " Balancing knowledge protection and knowledge sharing\n Knowledge protection risks", personas="1. Historian of Knowledge Systems: This editor will focus on the history and evolution of knowledge curation. They will provide context on how knowledge curation has changed over time and its impact on modern practices.\n" "2. Information Science Professional: With insights from 'Information science', this editor will explore the foundational theories, definitions, and philosophy that underpin knowledge curation\n" "3. Digital Librarian: This editor will delve into the specifics of how digital libraries operate, including software, metadata, digital preservation.\n" "4. Technical expert: This editor will focus on the technical aspects of knowledge curation, such as common features of content management systems.\n" "5. Museum Curator: The museum curator will contribute expertise on the curation of physical items and the transition of these practices into the digital realm.", ) runner.storm_knowledge_curation_module.persona_generator.create_writer_with_persona.find_related_topic.demos = [ find_related_topic_example ] runner.storm_knowledge_curation_module.persona_generator.create_writer_with_persona.gen_persona.demos = [ gen_persona_example ] # A trade-off of adding one-shot example is that it will increase the input length of the prompt. Also, some # examples may be very long (e.g., an example for writing a section based on the given information), which may # confuse the model. For these cases, you can create a pseudo-example that is short and easy to understand to steer # the model's output format. # For example, we can add the following pseudo-examples to the prompt used in WritePageOutlineFromConv and # ConvToSection. write_page_outline_example = Example( topic="Example Topic", conv="Wikipedia Writer: ...\nExpert: ...\nWikipedia Writer: ...\nExpert: ...", old_outline="# Section 1\n## Subsection 1\n## Subsection 2\n" "# Section 2\n## Subsection 1\n## Subsection 2\n" "# Section 3", outline="# New Section 1\n## New Subsection 1\n## New Subsection 2\n" "# New Section 2\n" "# New Section 3\n## New Subsection 1\n## New Subsection 2\n## New Subsection 3", ) runner.storm_outline_generation_module.write_outline.write_page_outline.demos = [ write_page_outline_example ] write_section_example = Example( info="[1]\nInformation in document 1\n[2]\nInformation in document 2\n[3]\nInformation in document 3", topic="Example Topic", section="Example Section", output="# Example Topic\n## Subsection 1\n" "This is an example sentence [1]. This is another example sentence [2][3].\n" "## Subsection 2\nThis is one more example sentence [1].", ) runner.storm_article_generation.section_gen.write_section.demos = [ write_section_example ] topic = input("Topic: ") runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--url", type=str, default="http://localhost", help="URL of the Ollama server." ) parser.add_argument( "--port", type=int, default=11434, help="Port of the Ollama server." ) parser.add_argument( "--model", type=str, default="llama3:latest", help="Model of the Ollama server." ) parser.add_argument( "--output-dir", type=str, default="./results/ollama", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "brave", "serper", "duckduckgo", "tavily", "searxng"], help="The search engine API to use for retrieving information.", ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_ollama_with_searxng.py ================================================ import os from argparse import ArgumentParser from dspy import Example from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import OllamaClient from knowledge_storm.rm import SearXNG from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() ollama_kwargs = { "model": args.model, "port": args.port, "url": args.url, "stop": ("\n\n---",), } conv_simulator_lm = OllamaClient(max_tokens=500, **ollama_kwargs) question_asker_lm = OllamaClient(max_tokens=500, **ollama_kwargs) outline_gen_lm = OllamaClient(max_tokens=400, **ollama_kwargs) article_gen_lm = OllamaClient(max_tokens=700, **ollama_kwargs) article_polish_lm = OllamaClient(max_tokens=4000, **ollama_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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) rm = SearXNG( searxng_api_url=args.searxng_api_url, searxng_api_key=os.getenv("SEARXNG_API_KEY"), k=engine_args.search_top_k, ) runner = STORMWikiRunner(engine_args, lm_configs, rm) find_related_topic_example = Example( topic="Knowledge Curation", related_topics="https://en.wikipedia.org/wiki/Knowledge_management\n" "https://en.wikipedia.org/wiki/Information_science\n" "https://en.wikipedia.org/wiki/Library_science\n", ) gen_persona_example = Example( topic="Knowledge Curation", examples="Title: Knowledge management\n" "Table of Contents: History\nResearch\n Dimensions\n Strategies\n Motivations\nKM technologies" "\nKnowledge barriers\nKnowledge retention\nKnowledge audit\nKnowledge protection\n" " Knowledge protection methods\n Formal methods\n Informal methods\n" " Balancing knowledge protection and knowledge sharing\n Knowledge protection risks", personas=( "1. Historian of Knowledge Systems: This editor will focus on the history and evolution of knowledge " "curation. They will provide context on how knowledge curation has changed over time and its impact on " "modern practices.\n" "2. Information Science Professional: With insights from 'Information science', this editor will " "explore the foundational theories, definitions, and philosophy that underpin knowledge curation\n" "3. Digital Librarian: This editor will delve into the specifics of how digital libraries operate, " "including software, metadata, digital preservation.\n" "4. Technical expert: This editor will focus on the technical aspects of knowledge curation, " "such as common features of content management systems.\n" "5. Museum Curator: The museum curator will contribute expertise on the curation of physical items and " "the transition of these practices into the digital realm." ), ) runner.storm_knowledge_curation_module.persona_generator.create_writer_with_persona.find_related_topic.demos = [ find_related_topic_example ] runner.storm_knowledge_curation_module.persona_generator.create_writer_with_persona.gen_persona.demos = [ gen_persona_example ] write_page_outline_example = Example( topic="Example Topic", conv="Wikipedia Writer: ...\nExpert: ...\nWikipedia Writer: ...\nExpert: ...", old_outline="# Section 1\n## Subsection 1\n## Subsection 2\n" "# Section 2\n## Subsection 1\n## Subsection 2\n" "# Section 3", outline="# New Section 1\n## New Subsection 1\n## New Subsection 2\n" "# New Section 2\n" "# New Section 3\n## New Subsection 1\n## New Subsection 2\n## New Subsection 3", ) runner.storm_outline_generation_module.write_outline.write_page_outline.demos = [ write_page_outline_example ] write_section_example = Example( info="[1]\nInformation in document 1\n[2]\nInformation in document 2\n[3]\nInformation in document 3", topic="Example Topic", section="Example Section", output="# Example Topic\n## Subsection 1\n" "This is an example sentence [1]. This is another example sentence [2][3].\n" "## Subsection 2\nThis is one more example sentence [1].", ) runner.storm_article_generation.section_gen.write_section.demos = [ write_section_example ] topic = input("Topic: ") runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--url", type=str, default="http://localhost", help="URL of the Ollama server." ) parser.add_argument( "--port", type=int, default=11434, help="Port of the Ollama server." ) parser.add_argument( "--model", type=str, default="llama3:latest", help="Model of the Ollama server." ) parser.add_argument( "--output-dir", type=str, default="./results/ollama", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["searxng"], help="The search engine API to use for retrieving information.", ) parser.add_argument( "--searxng-api-url", type=str, required=True, help="URL of the SearXNG API." ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: examples/storm_examples/run_storm_wiki_serper.py ================================================ """ STORM Wiki pipeline powered by Claude family models and serper search engine. You need to set up the following environment variables to run this script: - ANTHROPIC_API_KEY: Anthropic API key - SERPER_API_KEY: Serper.dev api key Output will be structured as below args.output_dir/ topic_name/ # topic_name will follow convention of underscore-connected topic name w/o space and slash conversation_log.json # Log of information-seeking conversation raw_search_results.json # Raw search results from search engine direct_gen_outline.txt # Outline directly generated with LLM's parametric knowledge storm_gen_outline.txt # Outline refined with collected information url_to_info.json # Sources that are used in the final article storm_gen_article.txt # Final article generated storm_gen_article_polished.txt # Polished final article (if args.do_polish_article is True) """ import os from argparse import ArgumentParser from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import ClaudeModel from knowledge_storm.rm import SerperRM from knowledge_storm.utils import load_api_key def main(args): load_api_key(toml_file_path="secrets.toml") lm_configs = STORMWikiLMConfigs() claude_kwargs = { "api_key": os.getenv("ANTHROPIC_API_KEY"), "temperature": 1.0, "top_p": 0.9, } # STORM is a LM system so different components can be powered by different models. # For a good balance between cost and quality, you can choose a cheaper/faster model for conv_simulator_lm # which is used to split queries, synthesize answers in the conversation. We recommend using stronger models # for outline_gen_lm which is responsible for organizing the collected information, and article_gen_lm # which is responsible for generating sections with citations. conv_simulator_lm = ClaudeModel( model="claude-3-haiku-20240307", max_tokens=500, **claude_kwargs ) question_asker_lm = ClaudeModel( model="claude-3-sonnet-20240229", max_tokens=500, **claude_kwargs ) outline_gen_lm = ClaudeModel( model="claude-3-opus-20240229", max_tokens=400, **claude_kwargs ) article_gen_lm = ClaudeModel( model="claude-3-opus-20240229", max_tokens=700, **claude_kwargs ) article_polish_lm = ClaudeModel( model="claude-3-opus-20240229", max_tokens=4000, **claude_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) engine_args = STORMWikiRunnerArguments( output_dir=args.output_dir, max_conv_turn=args.max_conv_turn, max_perspective=args.max_perspective, search_top_k=args.search_top_k, max_thread_num=args.max_thread_num, ) # Documentation to generate the data is available here: # https://serper.dev/playground # Important to note that tbs(date range is hardcoded values). # num is results per pages and is recommended to use in increments of 10(10, 20, etc). # page is how many pages will be searched. # h1 is where the google search will orginate from. topic = input("topic: ") data = {"autocorrect": True, "num": 10, "page": 1} rm = SerperRM(serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params=data) runner = STORMWikiRunner(engine_args, lm_configs, rm) runner.run( topic=topic, do_research=args.do_research, do_generate_outline=args.do_generate_outline, do_generate_article=args.do_generate_article, do_polish_article=args.do_polish_article, ) runner.post_run() runner.summary() if __name__ == "__main__": parser = ArgumentParser() # global arguments parser.add_argument( "--output-dir", type=str, default="./results/serper", help="Directory to store the outputs.", ) parser.add_argument( "--max-thread-num", type=int, default=3, help="Maximum number of threads to use. The information seeking part and the article generation" "part can speed up by using multiple threads. Consider reducing it if keep getting " '"Exceed rate limit" error when calling LM API.', ) parser.add_argument( "--retriever", type=str, choices=["bing", "you", "serper"], help="The search engine API to use for retrieving information.", ) # stage of the pipeline parser.add_argument( "--do-research", action="store_true", help="If True, simulate conversation to research the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-outline", action="store_true", help="If True, generate an outline for the topic; otherwise, load the results.", ) parser.add_argument( "--do-generate-article", action="store_true", help="If True, generate an article for the topic; otherwise, load the results.", ) parser.add_argument( "--do-polish-article", action="store_true", help="If True, polish the article by adding a summarization section and (optionally) removing " "duplicate content.", ) # hyperparameters for the pre-writing stage parser.add_argument( "--max-conv-turn", type=int, default=3, help="Maximum number of questions in conversational question asking.", ) parser.add_argument( "--max-perspective", type=int, default=3, help="Maximum number of perspectives to consider in perspective-guided question asking.", ) parser.add_argument( "--search-top-k", type=int, default=3, help="Top k search results to consider for each search query.", ) # hyperparameters for the writing stage parser.add_argument( "--retrieve-top-k", type=int, default=3, help="Top k collected references for each section title.", ) parser.add_argument( "--remove-duplicate", action="store_true", help="If True, remove duplicate content from the article.", ) main(parser.parse_args()) ================================================ FILE: frontend/demo_light/.streamlit/config.toml ================================================ [client] showErrorDetails = false toolbarMode = "minimal" [theme] primaryColor = "#F63366" backgroundColor = "#FFFFFF" secondaryBackgroundColor = "#F0F2F6" textColor = "#262730" font = "sans serif" ================================================ FILE: frontend/demo_light/README.md ================================================ # STORM Minimal User Interface This is a minimal user interface for `STORMWikiRunner` which includes the following features: 1. Allowing user to create a new article through the "Create New Article" page. 2. Showing the intermediate steps of STORMWikiRunner in real-time when creating an article. 3. Displaying the written article and references side by side. 4. Allowing user to view previously created articles through the "My Articles" page.

## Setup 1. Make sure you have installed `knowledge-storm` or set up the source code correctly. 2. Install additional packages required by the user interface: ```bash pip install -r requirements.txt ``` 2. Make sure you set up the API keys following the instructions in the main README file. Create a copy of `secrets.toml` and place it under `.streamlit/`. 3. Run the following command to start the user interface: ```bash streamlit run storm.py ``` The user interface will create a `DEMO_WORKING_DIR` directory in the current directory to store the outputs. ## Customization You can customize the `STORMWikiRunner` powering the user interface according to [the guidelines](https://github.com/stanford-oval/storm?tab=readme-ov-file#customize-storm) in the main README file. The `STORMWikiRunner` is initialized in `set_storm_runner()` in [demo_util.py](demo_util.py). You can change `STORMWikiRunnerArguments`, `STORMWikiLMConfigs`, or use a different retrieval model according to your need. ================================================ FILE: frontend/demo_light/demo_util.py ================================================ import base64 import datetime import json import os import re from typing import Optional import markdown import pytz import streamlit as st # If you install the source code instead of the `knowledge-storm` package, # Uncomment the following lines: # import sys # sys.path.append('../../') from knowledge_storm import ( STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs, ) from knowledge_storm.lm import OpenAIModel from knowledge_storm.rm import YouRM from knowledge_storm.storm_wiki.modules.callback import BaseCallbackHandler from knowledge_storm.utils import truncate_filename from stoc import stoc class DemoFileIOHelper: @staticmethod def read_structure_to_dict(articles_root_path): """ Reads the directory structure of articles stored in the given root path and returns a nested dictionary. The outer dictionary has article names as keys, and each value is another dictionary mapping file names to their absolute paths. Args: articles_root_path (str): The root directory path containing article subdirectories. Returns: dict: A dictionary where each key is an article name, and each value is a dictionary of file names and their absolute paths within that article's directory. """ articles_dict = {} for topic_name in os.listdir(articles_root_path): topic_path = os.path.join(articles_root_path, topic_name) if os.path.isdir(topic_path): # Initialize or update the dictionary for the topic articles_dict[topic_name] = {} # Iterate over all files within a topic directory for file_name in os.listdir(topic_path): file_path = os.path.join(topic_path, file_name) articles_dict[topic_name][file_name] = os.path.abspath(file_path) return articles_dict @staticmethod def read_txt_file(file_path): """ Reads the contents of a text file and returns it as a string. Args: file_path (str): The path to the text file to be read. Returns: str: The content of the file as a single string. """ with open(file_path) as f: return f.read() @staticmethod def read_json_file(file_path): """ Reads a JSON file and returns its content as a Python dictionary or list, depending on the JSON structure. Args: file_path (str): The path to the JSON file to be read. Returns: dict or list: The content of the JSON file. The type depends on the structure of the JSON file (object or array at the root). """ with open(file_path) as f: return json.load(f) @staticmethod def read_image_as_base64(image_path): """ Reads an image file and returns its content encoded as a base64 string, suitable for embedding in HTML or transferring over networks where binary data cannot be easily sent. Args: image_path (str): The path to the image file to be encoded. Returns: str: The base64 encoded string of the image, prefixed with the necessary data URI scheme for images. """ with open(image_path, "rb") as f: data = f.read() encoded = base64.b64encode(data) data = "data:image/png;base64," + encoded.decode("utf-8") return data @staticmethod def set_file_modification_time(file_path, modification_time_string): """ Sets the modification time of a file based on a given time string in the California time zone. Args: file_path (str): The path to the file. modification_time_string (str): The desired modification time in 'YYYY-MM-DD HH:MM:SS' format. """ california_tz = pytz.timezone("America/Los_Angeles") modification_time = datetime.datetime.strptime( modification_time_string, "%Y-%m-%d %H:%M:%S" ) modification_time = california_tz.localize(modification_time) modification_time_utc = modification_time.astimezone(datetime.timezone.utc) modification_timestamp = modification_time_utc.timestamp() os.utime(file_path, (modification_timestamp, modification_timestamp)) @staticmethod def get_latest_modification_time(path): """ Returns the latest modification time of all files in a directory in the California time zone as a string. Args: directory_path (str): The path to the directory. Returns: str: The latest file's modification time in 'YYYY-MM-DD HH:MM:SS' format. """ california_tz = pytz.timezone("America/Los_Angeles") latest_mod_time = None file_paths = [] if os.path.isdir(path): for root, dirs, files in os.walk(path): for file in files: file_paths.append(os.path.join(root, file)) else: file_paths = [path] for file_path in file_paths: modification_timestamp = os.path.getmtime(file_path) modification_time_utc = datetime.datetime.utcfromtimestamp( modification_timestamp ) modification_time_utc = modification_time_utc.replace( tzinfo=datetime.timezone.utc ) modification_time_california = modification_time_utc.astimezone( california_tz ) if ( latest_mod_time is None or modification_time_california > latest_mod_time ): latest_mod_time = modification_time_california if latest_mod_time is not None: return latest_mod_time.strftime("%Y-%m-%d %H:%M:%S") else: return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") @staticmethod def assemble_article_data(article_file_path_dict): """ Constructs a dictionary containing the content and metadata of an article based on the available files in the article's directory. This includes the main article text, citations from a JSON file, and a conversation log if available. The function prioritizes a polished version of the article if both a raw and polished version exist. Args: article_file_paths (dict): A dictionary where keys are file names relevant to the article (e.g., the article text, citations in JSON format, conversation logs) and values are their corresponding file paths. Returns: dict or None: A dictionary containing the parsed content of the article, citations, and conversation log if available. Returns None if neither the raw nor polished article text exists in the provided file paths. """ if ( "storm_gen_article.txt" in article_file_path_dict or "storm_gen_article_polished.txt" in article_file_path_dict ): full_article_name = ( "storm_gen_article_polished.txt" if "storm_gen_article_polished.txt" in article_file_path_dict else "storm_gen_article.txt" ) article_data = { "article": DemoTextProcessingHelper.parse( DemoFileIOHelper.read_txt_file( article_file_path_dict[full_article_name] ) ) } if "url_to_info.json" in article_file_path_dict: article_data["citations"] = _construct_citation_dict_from_search_result( DemoFileIOHelper.read_json_file( article_file_path_dict["url_to_info.json"] ) ) if "conversation_log.json" in article_file_path_dict: article_data["conversation_log"] = DemoFileIOHelper.read_json_file( article_file_path_dict["conversation_log.json"] ) return article_data return None class DemoTextProcessingHelper: @staticmethod def remove_citations(sent): return ( re.sub(r"\[\d+", "", re.sub(r" \[\d+", "", sent)) .replace(" |", "") .replace("]", "") ) @staticmethod def parse_conversation_history(json_data): """ Given conversation log data, return list of parsed data of following format (persona_name, persona_description, list of dialogue turn) """ parsed_data = [] for persona_conversation_data in json_data: if ": " in persona_conversation_data["perspective"]: name, description = persona_conversation_data["perspective"].split( ": ", 1 ) elif "- " in persona_conversation_data["perspective"]: name, description = persona_conversation_data["perspective"].split( "- ", 1 ) else: name, description = "", persona_conversation_data["perspective"] cur_conversation = [] for dialogue_turn in persona_conversation_data["dlg_turns"]: cur_conversation.append( {"role": "user", "content": dialogue_turn["user_utterance"]} ) cur_conversation.append( { "role": "assistant", "content": DemoTextProcessingHelper.remove_citations( dialogue_turn["agent_utterance"] ), } ) parsed_data.append((name, description, cur_conversation)) return parsed_data @staticmethod def parse(text): regex = re.compile(r']:\s+"(.*?)"\s+http') text = regex.sub("]: http", text) return text @staticmethod def add_markdown_indentation(input_string): lines = input_string.split("\n") processed_lines = [""] for line in lines: num_hashes = 0 for char in line: if char == "#": num_hashes += 1 else: break num_hashes -= 1 num_spaces = 4 * num_hashes new_line = " " * num_spaces + line processed_lines.append(new_line) return "\n".join(processed_lines) @staticmethod def get_current_time_string(): """ Returns the current time in the California time zone as a string. Returns: str: The current California time in 'YYYY-MM-DD HH:MM:SS' format. """ california_tz = pytz.timezone("America/Los_Angeles") utc_now = datetime.datetime.now(datetime.timezone.utc) california_now = utc_now.astimezone(california_tz) return california_now.strftime("%Y-%m-%d %H:%M:%S") @staticmethod def compare_time_strings( time_string1, time_string2, time_format="%Y-%m-%d %H:%M:%S" ): """ Compares two time strings to determine if they represent the same point in time. Args: time_string1 (str): The first time string to compare. time_string2 (str): The second time string to compare. time_format (str): The format of the time strings, defaults to '%Y-%m-%d %H:%M:%S'. Returns: bool: True if the time strings represent the same time, False otherwise. """ # Parse the time strings into datetime objects time1 = datetime.datetime.strptime(time_string1, time_format) time2 = datetime.datetime.strptime(time_string2, time_format) # Compare the datetime objects return time1 == time2 @staticmethod def add_inline_citation_link(article_text, citation_dict): # Regular expression to find citations like [i] pattern = r"\[(\d+)\]" # Function to replace each citation with its Markdown link def replace_with_link(match): i = match.group(1) url = citation_dict.get(int(i), {}).get("url", "#") return f"[[{i}]]({url})" # Replace all citations in the text with Markdown links return re.sub(pattern, replace_with_link, article_text) @staticmethod def generate_html_toc(md_text): toc = [] for line in md_text.splitlines(): if line.startswith("#"): level = line.count("#") title = line.strip("# ").strip() anchor = title.lower().replace(" ", "-").replace(".", "") toc.append( f"
  • {title}
    • " ) return "" @staticmethod def construct_bibliography_from_url_to_info(url_to_info): bibliography_list = [] sorted_url_to_unified_index = dict( sorted( url_to_info["url_to_unified_index"].items(), key=lambda item: item[1] ) ) for url, index in sorted_url_to_unified_index.items(): title = url_to_info["url_to_info"][url]["title"] bibliography_list.append(f"[{index}]: [{title}]({url})") bibliography_string = "\n\n".join(bibliography_list) return f"# References\n\n{bibliography_string}" class DemoUIHelper: def st_markdown_adjust_size(content, font_size=20): st.markdown( f""" {content} """, unsafe_allow_html=True, ) @staticmethod def get_article_card_UI_style(boarder_color="#9AD8E1"): return { "card": { "width": "100%", "height": "116px", "max-width": "640px", "background-color": "#FFFFF", "border": "1px solid #CCC", "padding": "20px", "border-radius": "5px", "border-left": f"0.5rem solid {boarder_color}", "box-shadow": "0 0.15rem 1.75rem 0 rgba(58, 59, 69, 0.15)", "margin": "0px", }, "title": { "white-space": "nowrap", "overflow": "hidden", "text-overflow": "ellipsis", "font-size": "17px", "color": "rgb(49, 51, 63)", "text-align": "left", "width": "95%", "font-weight": "normal", }, "text": { "white-space": "nowrap", "overflow": "hidden", "text-overflow": "ellipsis", "font-size": "25px", "color": "rgb(49, 51, 63)", "text-align": "left", "width": "95%", }, "filter": {"background-color": "rgba(0, 0, 0, 0)"}, } @staticmethod def customize_toast_css_style(): # Note padding is top right bottom left st.markdown( """ """, unsafe_allow_html=True, ) @staticmethod def article_markdown_to_html(article_title, article_content): return f""" {article_title}

    {article_title.replace('_', ' ')}

    Table of Contents

    {DemoTextProcessingHelper.generate_html_toc(article_content)} {markdown.markdown(article_content)} """ def _construct_citation_dict_from_search_result(search_results): if search_results is None: return None citation_dict = {} for url, index in search_results["url_to_unified_index"].items(): citation_dict[index] = { "url": url, "title": search_results["url_to_info"][url]["title"], "snippets": search_results["url_to_info"][url]["snippets"], } return citation_dict def _display_main_article_text(article_text, citation_dict, table_content_sidebar): # Post-process the generated article for better display. if "Write the lead section:" in article_text: article_text = article_text[ article_text.find("Write the lead section:") + len("Write the lead section:") : ] if article_text[0] == "#": article_text = "\n".join(article_text.split("\n")[1:]) article_text = DemoTextProcessingHelper.add_inline_citation_link( article_text, citation_dict ) # '$' needs to be changed to '\$' to avoid being interpreted as LaTeX in st.markdown() article_text = article_text.replace("$", "\\$") stoc.from_markdown(article_text, table_content_sidebar) def _display_references(citation_dict): if citation_dict: reference_list = [f"reference [{i}]" for i in range(1, len(citation_dict) + 1)] selected_key = st.selectbox("Select a reference", reference_list) citation_val = citation_dict[reference_list.index(selected_key) + 1] citation_val["title"] = citation_val["title"].replace("$", "\\$") st.markdown(f"**Title:** {citation_val['title']}") st.markdown(f"**Url:** {citation_val['url']}") snippets = "\n\n".join(citation_val["snippets"]).replace("$", "\\$") st.markdown(f"**Highlights:**\n\n {snippets}") else: st.markdown("**No references available**") def _display_persona_conversations(conversation_log): """ Display persona conversation in dialogue UI """ # get personas list as (persona_name, persona_description, dialogue turns list) tuple parsed_conversation_history = DemoTextProcessingHelper.parse_conversation_history( conversation_log ) # construct tabs for each persona conversation persona_tabs = st.tabs([name for (name, _, _) in parsed_conversation_history]) for idx, persona_tab in enumerate(persona_tabs): with persona_tab: # show persona description st.info(parsed_conversation_history[idx][1]) # show user / agent utterance in dialogue UI for message in parsed_conversation_history[idx][2]: message["content"] = message["content"].replace("$", "\\$") with st.chat_message(message["role"]): if message["role"] == "user": st.markdown(f"**{message['content']}**") else: st.markdown(message["content"]) def _display_main_article( selected_article_file_path_dict, show_reference=True, show_conversation=True ): article_data = DemoFileIOHelper.assemble_article_data( selected_article_file_path_dict ) with st.container(height=1000, border=True): table_content_sidebar = st.sidebar.expander( "**Table of contents**", expanded=True ) _display_main_article_text( article_text=article_data.get("article", ""), citation_dict=article_data.get("citations", {}), table_content_sidebar=table_content_sidebar, ) # display reference panel if show_reference and "citations" in article_data: with st.sidebar.expander("**References**", expanded=True): with st.container(height=800, border=False): _display_references(citation_dict=article_data.get("citations", {})) # display conversation history if show_conversation and "conversation_log" in article_data: with st.expander( "**STORM** is powered by a knowledge agent that proactively research a given topic by asking good questions coming from different perspectives.\n\n" ":sunglasses: Click here to view the agent's brain**STORM**ing process!" ): _display_persona_conversations( conversation_log=article_data.get("conversation_log", {}) ) def get_demo_dir(): return os.path.dirname(os.path.abspath(__file__)) def clear_other_page_session_state(page_index: Optional[int]): if page_index is None: keys_to_delete = [key for key in st.session_state if key.startswith("page")] else: keys_to_delete = [ key for key in st.session_state if key.startswith("page") and f"page{page_index}" not in key ] for key in set(keys_to_delete): del st.session_state[key] def set_storm_runner(): current_working_dir = os.path.join(get_demo_dir(), "DEMO_WORKING_DIR") if not os.path.exists(current_working_dir): os.makedirs(current_working_dir) # configure STORM runner llm_configs = STORMWikiLMConfigs() llm_configs.init_openai_model( openai_api_key=st.secrets["OPENAI_API_KEY"], openai_type="openai" ) llm_configs.set_question_asker_lm( OpenAIModel( model="gpt-4-1106-preview", api_key=st.secrets["OPENAI_API_KEY"], api_provider="openai", max_tokens=500, temperature=1.0, top_p=0.9, ) ) engine_args = STORMWikiRunnerArguments( output_dir=current_working_dir, max_conv_turn=3, max_perspective=3, search_top_k=3, retrieve_top_k=5, ) rm = YouRM(ydc_api_key=st.secrets["YDC_API_KEY"], k=engine_args.search_top_k) runner = STORMWikiRunner(engine_args, llm_configs, rm) st.session_state["runner"] = runner def display_article_page( selected_article_name, selected_article_file_path_dict, show_title=True, show_main_article=True, ): if show_title: st.markdown( f"

    {selected_article_name.replace('_', ' ')}

    ", unsafe_allow_html=True, ) if show_main_article: _display_main_article(selected_article_file_path_dict) class StreamlitCallbackHandler(BaseCallbackHandler): def __init__(self, status_container): self.status_container = status_container def on_identify_perspective_start(self, **kwargs): self.status_container.info( "Start identifying different perspectives for researching the topic." ) def on_identify_perspective_end(self, perspectives: list[str], **kwargs): perspective_list = "\n- ".join(perspectives) self.status_container.success( f"Finish identifying perspectives. Will now start gathering information" f" from the following perspectives:\n- {perspective_list}" ) def on_information_gathering_start(self, **kwargs): self.status_container.info("Start browsing the Internet.") def on_dialogue_turn_end(self, dlg_turn, **kwargs): urls = list(set([r.url for r in dlg_turn.search_results])) for url in urls: self.status_container.markdown( f"""
    Finish browsing {url}.
    """, unsafe_allow_html=True, ) def on_information_gathering_end(self, **kwargs): self.status_container.success("Finish collecting information.") def on_information_organization_start(self, **kwargs): self.status_container.info( "Start organizing information into a hierarchical outline." ) def on_direct_outline_generation_end(self, outline: str, **kwargs): self.status_container.success( f"Finish leveraging the internal knowledge of the large language model." ) def on_outline_refinement_end(self, outline: str, **kwargs): self.status_container.success(f"Finish leveraging the collected information.") ================================================ FILE: frontend/demo_light/pages_util/CreateNewArticle.py ================================================ import os import time import demo_util import streamlit as st from demo_util import ( DemoFileIOHelper, DemoTextProcessingHelper, DemoUIHelper, truncate_filename, ) def handle_not_started(): if st.session_state["page3_write_article_state"] == "not started": _, search_form_column, _ = st.columns([2, 5, 2]) with search_form_column: with st.form(key="search_form"): # Text input for the search topic DemoUIHelper.st_markdown_adjust_size( content="Enter the topic you want to learn in depth:", font_size=18 ) st.session_state["page3_topic"] = st.text_input( label="page3_topic", label_visibility="collapsed" ) pass_appropriateness_check = True # Submit button for the form submit_button = st.form_submit_button(label="Research") # only start new search when button is clicked, not started, or already finished previous one if submit_button and st.session_state["page3_write_article_state"] in [ "not started", "show results", ]: if not st.session_state["page3_topic"].strip(): pass_appropriateness_check = False st.session_state["page3_warning_message"] = ( "topic could not be empty" ) st.session_state["page3_topic_name_cleaned"] = ( st.session_state["page3_topic"] .replace(" ", "_") .replace("/", "_") ) st.session_state["page3_topic_name_truncated"] = truncate_filename( st.session_state["page3_topic_name_cleaned"] ) if not pass_appropriateness_check: st.session_state["page3_write_article_state"] = "not started" alert = st.warning( st.session_state["page3_warning_message"], icon="⚠️" ) time.sleep(5) alert.empty() else: st.session_state["page3_write_article_state"] = "initiated" def handle_initiated(): if st.session_state["page3_write_article_state"] == "initiated": current_working_dir = os.path.join(demo_util.get_demo_dir(), "DEMO_WORKING_DIR") if not os.path.exists(current_working_dir): os.makedirs(current_working_dir) if "runner" not in st.session_state: demo_util.set_storm_runner() st.session_state["page3_current_working_dir"] = current_working_dir st.session_state["page3_write_article_state"] = "pre_writing" def handle_pre_writing(): if st.session_state["page3_write_article_state"] == "pre_writing": status = st.status( "I am brain**STORM**ing now to research the topic. (This may take 2-3 minutes.)" ) st_callback_handler = demo_util.StreamlitCallbackHandler(status) with status: # STORM main gen outline st.session_state["runner"].run( topic=st.session_state["page3_topic"], do_research=True, do_generate_outline=True, do_generate_article=False, do_polish_article=False, callback_handler=st_callback_handler, ) conversation_log_path = os.path.join( st.session_state["page3_current_working_dir"], st.session_state["page3_topic_name_truncated"], "conversation_log.json", ) demo_util._display_persona_conversations( DemoFileIOHelper.read_json_file(conversation_log_path) ) st.session_state["page3_write_article_state"] = "final_writing" status.update(label="brain**STORM**ing complete!", state="complete") def handle_final_writing(): if st.session_state["page3_write_article_state"] == "final_writing": # polish final article with st.status( "Now I will connect the information I found for your reference. (This may take 4-5 minutes.)" ) as status: st.info( "Now I will connect the information I found for your reference. (This may take 4-5 minutes.)" ) st.session_state["runner"].run( topic=st.session_state["page3_topic"], do_research=False, do_generate_outline=False, do_generate_article=True, do_polish_article=True, remove_duplicate=False, ) # finish the session st.session_state["runner"].post_run() # update status bar st.session_state["page3_write_article_state"] = "prepare_to_show_result" status.update(label="information snythesis complete!", state="complete") def handle_prepare_to_show_result(): if st.session_state["page3_write_article_state"] == "prepare_to_show_result": _, show_result_col, _ = st.columns([4, 3, 4]) with show_result_col: if st.button("show final article"): st.session_state["page3_write_article_state"] = "completed" st.rerun() def handle_completed(): if st.session_state["page3_write_article_state"] == "completed": # display polished article current_working_dir_paths = DemoFileIOHelper.read_structure_to_dict( st.session_state["page3_current_working_dir"] ) current_article_file_path_dict = current_working_dir_paths[ st.session_state["page3_topic_name_truncated"] ] demo_util.display_article_page( selected_article_name=st.session_state["page3_topic_name_cleaned"], selected_article_file_path_dict=current_article_file_path_dict, show_title=True, show_main_article=True, ) def create_new_article_page(): demo_util.clear_other_page_session_state(page_index=3) if "page3_write_article_state" not in st.session_state: st.session_state["page3_write_article_state"] = "not started" handle_not_started() handle_initiated() handle_pre_writing() handle_final_writing() handle_prepare_to_show_result() handle_completed() ================================================ FILE: frontend/demo_light/pages_util/MyArticles.py ================================================ import os import demo_util import streamlit as st from demo_util import DemoFileIOHelper, DemoUIHelper from streamlit_card import card # set page config and display title def my_articles_page(): with st.sidebar: _, return_button_col = st.columns([2, 5]) with return_button_col: if st.button( "Select another article", disabled="page2_selected_my_article" not in st.session_state, ): if "page2_selected_my_article" in st.session_state: del st.session_state["page2_selected_my_article"] st.rerun() # sync my articles if "page2_user_articles_file_path_dict" not in st.session_state: local_dir = os.path.join(demo_util.get_demo_dir(), "DEMO_WORKING_DIR") os.makedirs(local_dir, exist_ok=True) st.session_state["page2_user_articles_file_path_dict"] = ( DemoFileIOHelper.read_structure_to_dict(local_dir) ) # if no feature demo selected, display all featured articles as info cards def article_card_setup(column_to_add, card_title, article_name): with column_to_add: cleaned_article_title = article_name.replace("_", " ") hasClicked = card( title=" / ".join(card_title), text=article_name.replace("_", " "), image=DemoFileIOHelper.read_image_as_base64( os.path.join(demo_util.get_demo_dir(), "assets", "void.jpg") ), styles=DemoUIHelper.get_article_card_UI_style(boarder_color="#9AD8E1"), ) if hasClicked: st.session_state["page2_selected_my_article"] = article_name st.rerun() if "page2_selected_my_article" not in st.session_state: # display article cards my_article_columns = st.columns(3) if len(st.session_state["page2_user_articles_file_path_dict"]) > 0: # get article names article_names = sorted( list(st.session_state["page2_user_articles_file_path_dict"].keys()) ) # configure pagination pagination = st.container() bottom_menu = st.columns((1, 4, 1, 1, 1))[1:-1] with bottom_menu[2]: batch_size = st.selectbox("Page Size", options=[24, 48, 72]) with bottom_menu[1]: total_pages = ( int(len(article_names) / batch_size) if int(len(article_names) / batch_size) > 0 else 1 ) current_page = st.number_input( "Page", min_value=1, max_value=total_pages, step=1 ) with bottom_menu[0]: st.markdown(f"Page **{current_page}** of **{total_pages}** ") # show article cards with pagination: my_article_count = 0 start_index = (current_page - 1) * batch_size end_index = min(current_page * batch_size, len(article_names)) for article_name in article_names[start_index:end_index]: column_to_add = my_article_columns[my_article_count % 3] my_article_count += 1 article_card_setup( column_to_add=column_to_add, card_title=["My Article"], article_name=article_name, ) else: with my_article_columns[0]: hasClicked = card( title="Get started", text="Start your first research!", image=DemoFileIOHelper.read_image_as_base64( os.path.join(demo_util.get_demo_dir(), "assets", "void.jpg") ), styles=DemoUIHelper.get_article_card_UI_style(), ) if hasClicked: st.session_state.selected_page = 1 st.session_state["manual_selection_override"] = True st.session_state["rerun_requested"] = True st.rerun() else: selected_article_name = st.session_state["page2_selected_my_article"] selected_article_file_path_dict = st.session_state[ "page2_user_articles_file_path_dict" ][selected_article_name] demo_util.display_article_page( selected_article_name=selected_article_name, selected_article_file_path_dict=selected_article_file_path_dict, show_title=True, show_main_article=True, ) ================================================ FILE: frontend/demo_light/requirements.txt ================================================ streamlit==1.31.1 streamlit-card markdown unidecode extra-streamlit-components==0.1.60 streamlit_extras deprecation==2.1.0 st-pages==0.4.5 streamlit-float streamlit-option-menu ================================================ FILE: frontend/demo_light/stoc.py ================================================ """https://github.com/arnaudmiribel/stoc""" import re import streamlit as st import unidecode DISABLE_LINK_CSS = """ """ class stoc: def __init__(self): self.toc_items = list() def h1(self, text: str, write: bool = True): if write: st.write(f"# {text}") self.toc_items.append(("h1", text)) def h2(self, text: str, write: bool = True): if write: st.write(f"## {text}") self.toc_items.append(("h2", text)) def h3(self, text: str, write: bool = True): if write: st.write(f"### {text}") self.toc_items.append(("h3", text)) def toc(self, expander): st.write(DISABLE_LINK_CSS, unsafe_allow_html=True) # st.sidebar.caption("Table of contents") if expander is None: expander = st.sidebar.expander("**Table of contents**", expanded=True) with expander: with st.container(height=600, border=False): markdown_toc = "" for title_size, title in self.toc_items: h = int(title_size.replace("h", "")) markdown_toc += ( " " * 2 * h + "- " + f' {title} \n' ) # st.sidebar.write(markdown_toc, unsafe_allow_html=True) st.write(markdown_toc, unsafe_allow_html=True) @classmethod def get_toc(cls, markdown_text: str, topic=""): def increase_heading_depth_and_add_top_heading(markdown_text, new_top_heading): lines = markdown_text.splitlines() # Increase the depth of each heading by adding an extra '#' increased_depth_lines = [ "#" + line if line.startswith("#") else line for line in lines ] # Add the new top-level heading at the beginning increased_depth_lines.insert(0, f"# {new_top_heading}") # Re-join the modified lines back into a single string modified_text = "\n".join(increased_depth_lines) return modified_text if topic: markdown_text = increase_heading_depth_and_add_top_heading( markdown_text, topic ) toc = [] for line in markdown_text.splitlines(): if line.startswith("#"): # Remove the '#' characters and strip leading/trailing spaces heading_text = line.lstrip("#").strip() # Create slug (lowercase, spaces to hyphens, remove non-alphanumeric characters) slug = ( re.sub(r"[^a-zA-Z0-9\s-]", "", heading_text) .lower() .replace(" ", "-") ) # Determine heading level for indentation level = line.count("#") - 1 # Add to the table of contents toc.append(" " * level + f"- [{heading_text}](#{slug})") return "\n".join(toc) @classmethod def from_markdown(cls, text: str, expander=None): self = cls() for line in text.splitlines(): if line.startswith("###"): self.h3(line[3:], write=False) elif line.startswith("##"): self.h2(line[2:], write=False) elif line.startswith("#"): self.h1(line[1:], write=False) # customize markdown font size custom_css = """ """ st.markdown(custom_css, unsafe_allow_html=True) st.write(text) self.toc(expander=expander) def normalize(s): """ Normalize titles as valid HTML ids for anchors >>> normalize("it's a test to spot how Things happ3n héhé") "it-s-a-test-to-spot-how-things-happ3n-h-h" """ # Replace accents with "-" s_wo_accents = unidecode.unidecode(s) accents = [s for s in s if s not in s_wo_accents] for accent in accents: s = s.replace(accent, "-") # Lowercase s = s.lower() # Keep only alphanum and remove "-" suffix if existing normalized = ( "".join([char if char.isalnum() else "-" for char in s]).strip("-").lower() ) return normalized ================================================ FILE: frontend/demo_light/storm.py ================================================ import os script_dir = os.path.dirname(os.path.abspath(__file__)) wiki_root_dir = os.path.dirname(os.path.dirname(script_dir)) import demo_util from pages_util import MyArticles, CreateNewArticle from streamlit_float import * from streamlit_option_menu import option_menu def main(): global database st.set_page_config(layout="wide") if "first_run" not in st.session_state: st.session_state["first_run"] = True # set api keys from secrets if st.session_state["first_run"]: for key, value in st.secrets.items(): if type(value) == str: os.environ[key] = value # initialize session_state if "selected_article_index" not in st.session_state: st.session_state["selected_article_index"] = 0 if "selected_page" not in st.session_state: st.session_state["selected_page"] = 0 if st.session_state.get("rerun_requested", False): st.session_state["rerun_requested"] = False st.rerun() st.write( "", unsafe_allow_html=True ) menu_container = st.container() with menu_container: pages = ["My Articles", "Create New Article"] styles = { "container": {"padding": "0.2rem 0", "background-color": "#22222200"}, } menu_selection = option_menu( None, pages, icons=["house", "search"], menu_icon="cast", default_index=0, orientation="horizontal", manual_select=st.session_state.selected_page, styles=styles, key="menu_selection", ) if st.session_state.get("manual_selection_override", False): menu_selection = pages[st.session_state["selected_page"]] st.session_state["manual_selection_override"] = False st.session_state["selected_page"] = None if menu_selection == "My Articles": demo_util.clear_other_page_session_state(page_index=2) MyArticles.my_articles_page() elif menu_selection == "Create New Article": demo_util.clear_other_page_session_state(page_index=3) CreateNewArticle.create_new_article_page() if __name__ == "__main__": main() ================================================ FILE: knowledge_storm/__init__.py ================================================ from .storm_wiki import * from .collaborative_storm import * from .encoder import * from .interface import * from .lm import * from .rm import * from .utils import * from .dataclass import * __version__ = "1.1.0" ================================================ FILE: knowledge_storm/collaborative_storm/__init__.py ================================================ from .modules import * from .engine import * ================================================ FILE: knowledge_storm/collaborative_storm/engine.py ================================================ import dspy import os from dataclasses import dataclass, field, asdict from typing import List, Union, Literal, Optional, Dict from .modules import collaborative_storm_utils as collaborative_storm_utils from .modules.callback import BaseCallbackHandler from .modules.co_storm_agents import ( SimulatedUser, PureRAGAgent, Moderator, CoStormExpert, ) from .modules.expert_generation import GenerateExpertModule from .modules.warmstart_hierarchical_chat import WarmStartModule from ..dataclass import ConversationTurn, KnowledgeBase from ..encoder import Encoder from ..interface import LMConfigs, Agent from ..logging_wrapper import LoggingWrapper from ..lm import LitellmModel from ..rm import BingSearch class CollaborativeStormLMConfigs(LMConfigs): """Configurations for LLM used in different parts of Co-STORM. Given that different parts in Co-STORM framework have different complexity, we use different LLM configurations to achieve a balance between quality and efficiency. If no specific configuration is provided, we use the default setup in the paper. """ def __init__(self): self.question_answering_lm = None self.discourse_manage_lm = None self.utterance_polishing_lm = None self.warmstart_outline_gen_lm = None self.question_asking_lm = None self.knowledge_base_lm = None def init( self, lm_type: Literal["openai", "azure", "together"], temperature: Optional[float] = 1.0, top_p: Optional[float] = 0.9, ): if lm_type and lm_type == "openai": openai_kwargs = { "api_key": os.getenv("OPENAI_API_KEY"), "temperature": temperature, "top_p": top_p, "api_base": None, } self.question_answering_lm = LitellmModel( model="gpt-4o-2024-05-13", max_tokens=1000, **openai_kwargs ) self.discourse_manage_lm = LitellmModel( model="gpt-4o-2024-05-13", max_tokens=500, **openai_kwargs ) self.utterance_polishing_lm = LitellmModel( model="gpt-4o-2024-05-13", max_tokens=2000, **openai_kwargs ) self.warmstart_outline_gen_lm = LitellmModel( model="gpt-4-1106-preview", max_tokens=500, **openai_kwargs ) self.question_asking_lm = LitellmModel( model="gpt-4o-2024-05-13", max_tokens=300, **openai_kwargs ) self.knowledge_base_lm = LitellmModel( model="gpt-4o-2024-05-13", max_tokens=1000, **openai_kwargs ) elif lm_type and lm_type == "azure": azure_kwargs = { "api_key": os.getenv("AZURE_API_KEY"), "temperature": temperature, "top_p": top_p, "api_base": os.getenv("AZURE_API_BASE"), "api_version": os.getenv("AZURE_API_VERSION"), } self.question_answering_lm = LitellmModel( model="azure/gpt-4o", max_tokens=1000, **azure_kwargs, model_type="chat" ) self.discourse_manage_lm = LitellmModel( model="azure/gpt-4o", max_tokens=500, **azure_kwargs, model_type="chat" ) self.utterance_polishing_lm = LitellmModel( model="azure/gpt-4o", max_tokens=2000, **azure_kwargs, model_type="chat" ) self.warmstart_outline_gen_lm = LitellmModel( model="azure/gpt-4o", max_tokens=300, **azure_kwargs, model_type="chat" ) self.question_asking_lm = LitellmModel( model="azure/gpt-4o", max_tokens=300, **azure_kwargs, model_type="chat" ) self.knowledge_base_lm = LitellmModel( model="azure/gpt-4o", max_tokens=1000, **azure_kwargs, model_type="chat" ) elif lm_type and lm_type == "together": together_kwargs = { "api_key": os.getenv("TOGETHER_API_KEY"), "temperature": temperature, "top_p": top_p, } self.question_answering_lm = LitellmModel( model="together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max_tokens=1000, model_type="chat", **together_kwargs, ) self.discourse_manage_lm = LitellmModel( model="together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max_tokens=500, model_type="chat", **together_kwargs, ) self.utterance_polishing_lm = LitellmModel( model="together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max_tokens=2000, model_type="chat", **together_kwargs, ) self.warmstart_outline_gen_lm = LitellmModel( model="together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max_tokens=500, model_type="chat", **together_kwargs, ) self.question_asking_lm = LitellmModel( model="together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max_tokens=300, model_type="chat", **together_kwargs, ) self.knowledge_base_lm = LitellmModel( model="together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", max_tokens=1000, model_type="chat", **together_kwargs, ) else: raise Exception( "No valid OpenAI API provider is provided. Cannot use default LLM configurations." ) def set_question_answering_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.question_answering_lm = model def set_discourse_manage_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.discourse_manage_lm = model def set_utterance_polishing_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.utterance_polishing_lm = model def set_warmstart_outline_gen_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.warmstart_outline_gen_lm = model def set_question_asking_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.question_asking_lm = model def set_knowledge_base_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.knowledge_base_lm = model def collect_and_reset_lm_usage(self): lm_usage = {} for attr_name in self.__dict__: if "_lm" in attr_name and hasattr( getattr(self, attr_name), "get_usage_and_reset" ): usage = getattr(self, attr_name).get_usage_and_reset() if any( value["prompt_tokens"] != 0 or value["completion_tokens"] != 0 for value in usage.values() ): lm_usage[attr_name] = usage return lm_usage def to_dict(self): """ Converts the CollaborativeStormLMConfigs instance to a dictionary representation. Returns: dict: The dictionary representation of the CollaborativeStormLMConfigs. """ config_dict = {} for attr_name in self.__dict__: config_dict[attr_name] = getattr(self, attr_name).kwargs return config_dict @dataclass class RunnerArgument: """Arguments for controlling the STORM Wiki pipeline.""" topic: str = field( metadata={"help": "Topic of discourse"}, ) retrieve_top_k: int = field( default=10, metadata={"help": "retrieve top k results for each query in retriever"}, ) max_search_queries: int = field( default=2, metadata={ "help": "Maximum number of search queries to consider for each question." }, ) total_conv_turn: int = field( default=20, metadata={"help": "Maximum number turn in conversation."}, ) max_search_thread: int = field( default=5, metadata={"help": "Maximum number of parallel thread for retriever"}, ) max_search_queries_per_turn: int = field( default=3, metadata={"help": "Maximum number of search queries to consider in each turn."}, ) warmstart_max_num_experts: int = field( default=3, metadata={ "help": "Max number of experts in perspective guided QA in warm start process" }, ) warmstart_max_turn_per_experts: int = field( default=2, metadata={"help": "Max number of turns per perspective in warm start process"}, ) warmstart_max_thread: int = field( default=3, metadata={ "help": "Max number thread for parallel perspective guided QA in warm start process" }, ) max_thread_num: int = field( default=10, metadata={ "help": "Maximum number of threads to use. " "Consider reducing it if keep getting 'Exceed rate limit' error when calling LM API." }, ) max_num_round_table_experts: int = field( default=2, metadata={"help": "Max number of active experts in round table discussion."}, ) moderator_override_N_consecutive_answering_turn: int = field( default=3, metadata={ "help": "Number of consecutive experts answering turn before moderator override the conversation" }, ) node_expansion_trigger_count: int = field( default=10, metadata={ "help": "Trigger node expansion for node that contain more than N snippets" }, ) disable_moderator: bool = field( default=False, metadata={"help": "If True, disable moderator."}, ) disable_multi_experts: bool = field( default=False, metadata={"help": "If True, disable moderator."}, ) rag_only_baseline_mode: bool = field( default=False, metadata={"help": "If True, switch to rag online baseline mode"}, ) def to_dict(self): """ Converts the RunnerArgument instance to a dictionary representation. Returns: dict: The dictionary representation of the RunnerArgument. """ return asdict(self) @classmethod def from_dict(cls, data): """ Constructs a RunnerArgument instance from a dictionary representation. Args: data (dict): The dictionary representation of the RunnerArgument. Returns: RunnerArgument: The constructed RunnerArgument instance. """ return cls(**data) @dataclass class TurnPolicySpec: """ Represents the policy specifications for determining the behavior of a conversation turn. Attributes: should_reorganize_knowledge_base (bool): A flag that indicates whether the knowledge base should be reorganized after the current turn. should_update_experts_list (bool): A flag that indicates whether the list of experts should be updated based on the conversation context. should_polish_utterance (bool): A flag that indicates whether the generated utterance should be polished (e.g., refined or rephrased) before it is used in the conversation. agent (Agent): The `Agent` responsible for generating utterances or responses during the conversation turn. This agent interacts with the knowledge base and the conversation history to produce responses. """ should_reorganize_knowledge_base: bool = False should_update_experts_list: bool = False should_polish_utterance: bool = False agent: Agent = None class DiscourseManager: def __init__( self, logging_wrapper: LoggingWrapper, lm_config: CollaborativeStormLMConfigs, runner_argument: RunnerArgument, rm: dspy.Retrieve, encoder: Encoder, callback_handler: BaseCallbackHandler, ): # parameter management self.lm_config = lm_config self.runner_argument = runner_argument self.logging_wrapper = logging_wrapper self.callback_handler = callback_handler self.rm = rm self.encoder = encoder # role management self.experts: List[CoStormExpert] = [] self.simulated_user: SimulatedUser = SimulatedUser( topic=self.runner_argument.topic, role_name="Guest", role_description="", intent=None, lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, callback_handler=self.callback_handler, ) self.pure_rag_agent: PureRAGAgent = PureRAGAgent( topic=self.runner_argument.topic, role_name="PureRAG", role_description="", lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=self.rm, callback_handler=self.callback_handler, ) self.moderator: Moderator = Moderator( topic=self.runner_argument.topic, role_name="Moderator", role_description="", lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, encoder=self.encoder, callback_handler=self.callback_handler, ) self.general_knowledge_provider = CoStormExpert( topic=self.runner_argument.topic, role_name="General Knowledge Provider", role_description="Focus on broadly covering the basic facts about the question.", lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=self.rm, callback_handler=self.callback_handler, ) self.generate_expert_module = GenerateExpertModule( engine=self.lm_config.discourse_manage_lm ) self.next_turn_moderator_override = False def serialize_experts(self) -> List[Dict]: return [ { "topic": expert.topic, "role_name": expert.role_name, "role_description": expert.role_description, } for expert in self.experts ] def deserialize_experts(self, data: List[Dict]): for expert_data in data: self.experts.append( CoStormExpert( topic=expert_data["topic"], role_name=expert_data["role_name"], role_description=expert_data["role_description"], lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=self.rm, callback_handler=self.callback_handler, ) ) def _should_generate_question( self, conversation_history: List[ConversationTurn] ) -> bool: consecutive_non_questioning_turn = 0 for conv_turn in reversed(conversation_history): if conv_turn.utterance_type not in [ "Original Question", "Information Request", ]: consecutive_non_questioning_turn += 1 else: break return ( consecutive_non_questioning_turn >= self.runner_argument.moderator_override_N_consecutive_answering_turn ) def _parse_expert_names_to_agent(self, expert_descriptions: Union[str, List[str]]): if type(expert_descriptions) == str: expert_descriptions = [expert_descriptions] agents: CoStormExpert = [] for expert_name in expert_descriptions: role_name, role_description = expert_name.split(":") role_name = role_name.strip() role_description = role_description.strip() new_costorm_expert = CoStormExpert( topic=self.runner_argument.topic, role_name=role_name, role_description=role_description, lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=self.rm, callback_handler=self.callback_handler, ) agents.append(new_costorm_expert) return agents def _update_expert_list_from_utterance(self, focus: str, background_info: str): expert_names = self.generate_expert_module( topic=self.runner_argument.topic, background_info=background_info, focus=focus, num_experts=self.runner_argument.max_num_round_table_experts, ).experts self.experts = self._parse_expert_names_to_agent(expert_names) def _is_last_turn_questioning(self, conversation_history: List[ConversationTurn]): return conversation_history and conversation_history[-1].utterance_type in [ "Original Question", "Information Request", ] def get_next_turn_policy( self, conversation_history: List[ConversationTurn], dry_run=False, simulate_user=False, simulate_user_intent: str = None, ) -> TurnPolicySpec: next_turn_policy = TurnPolicySpec() if simulate_user: self.simulated_user.intent = simulate_user_intent next_turn_policy.agent = self.simulated_user elif self.runner_argument.rag_only_baseline_mode: assert self.conversation_history[-1].role == "Guest" next_turn_policy.agent = self.pure_rag_agent elif self.next_turn_moderator_override: next_turn_policy.agent = self.moderator if not dry_run: self.next_turn_moderator_override = False elif ( not self.runner_argument.disable_moderator and self._should_generate_question(conversation_history) ): next_turn_policy.agent = self.moderator next_turn_policy.should_reorganize_knowledge_base = True # experts RAG gen else: next_turn_policy.agent = self.general_knowledge_provider if ( not self._is_last_turn_questioning(conversation_history) and not self.runner_argument.disable_multi_experts ): if dry_run: next_turn_policy.agent = self.experts[0] else: next_turn_policy.agent = self.experts.pop(0) self.experts.append(next_turn_policy.agent) next_turn_policy.should_update_experts_list = ( self._is_last_turn_questioning(conversation_history) and not self.runner_argument.disable_multi_experts ) next_turn_policy.should_polish_utterance = True return next_turn_policy class CoStormRunner: def __init__( self, lm_config: CollaborativeStormLMConfigs, runner_argument: RunnerArgument, logging_wrapper: LoggingWrapper, rm: Optional[dspy.Retrieve] = None, callback_handler: BaseCallbackHandler = None, ): self.runner_argument = runner_argument self.lm_config = lm_config self.logging_wrapper = logging_wrapper self.callback_handler = callback_handler if rm is None: self.rm = BingSearch(k=runner_argument.retrieve_top_k) else: self.rm = rm self.encoder = Encoder() self.conversation_history = [] self.warmstart_conv_archive = [] self.knowledge_base = KnowledgeBase( topic=self.runner_argument.topic, knowledge_base_lm=self.lm_config.knowledge_base_lm, node_expansion_trigger_count=self.runner_argument.node_expansion_trigger_count, encoder=self.encoder, ) self.discourse_manager = DiscourseManager( lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=self.rm, encoder=self.encoder, callback_handler=callback_handler, ) def to_dict(self): return { "runner_argument": self.runner_argument.to_dict(), "lm_config": self.lm_config.to_dict(), "conversation_history": [ turn.to_dict() for turn in self.conversation_history ], "warmstart_conv_archive": [ turn.to_dict() for turn in self.warmstart_conv_archive ], "experts": self.discourse_manager.serialize_experts(), "knowledge_base": self.knowledge_base.to_dict(), } @classmethod def from_dict(cls, data, callback_handler: BaseCallbackHandler = None): # FIXME: does not use the lm_config data but naively use default setting lm_config = CollaborativeStormLMConfigs() lm_config.init(lm_type=os.getenv("OPENAI_API_TYPE")) costorm_runner = cls( lm_config=lm_config, runner_argument=RunnerArgument.from_dict(data["runner_argument"]), logging_wrapper=LoggingWrapper(lm_config), callback_handler=callback_handler, ) costorm_runner.encoder = Encoder() costorm_runner.conversation_history = [ ConversationTurn.from_dict(turn) for turn in data["conversation_history"] ] costorm_runner.warmstart_conv_archive = [ ConversationTurn.from_dict(turn) for turn in data.get("warmstart_conv_archive", []) ] costorm_runner.discourse_manager.deserialize_experts(data["experts"]) costorm_runner.knowledge_base = KnowledgeBase.from_dict( data=data["knowledge_base"], knowledge_base_lm=costorm_runner.lm_config.knowledge_base_lm, node_expansion_trigger_count=costorm_runner.runner_argument.node_expansion_trigger_count, encoder=costorm_runner.encoder, ) return costorm_runner def warm_start(self): """ Warm start co-storm system to conduct background information search in order to build shared conceptual space with user. This stage is a mini-STORM, spawning multiple LLM agent with different perspective and perform multi-round conversation. The knowledge base (i.e. mind map) will be initialize using the collected information. It will also generate a first draft of report and use it to produce an engaging and concise conversation presented to the user to catch up with system's knowledge about the topic. """ with self.logging_wrapper.log_pipeline_stage( pipeline_stage=f"warm start stage" ): if not self.runner_argument.rag_only_baseline_mode: warm_start_module = WarmStartModule( lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=self.rm, callback_handler=self.callback_handler, ) ( warmstart_conv, warmstart_revised_conv, warmstart_experts, ) = warm_start_module.initiate_warm_start( topic=self.runner_argument.topic, knowledge_base=self.knowledge_base, ) self.discourse_manager.experts = ( self.discourse_manager._parse_expert_names_to_agent( warmstart_experts ) ) self.discourse_manager.next_turn_moderator_override = True self.conversation_history = ( warmstart_revised_conv if warmstart_revised_conv else warmstart_conv ) self.warmstart_conv_archive = warmstart_conv self.knowledge_base.reorganize() else: if self.knowledge_base is None: self.knowledge_base = KnowledgeBase( topic=self.runner_argument.topic, knowledge_base_lm=self.lm_config.knowledge_base_lm, node_expansion_trigger_count=self.runner_argument.node_expansion_trigger_count, encoder=self.encoder, ) if self.conversation_history is None: self.conversation_history = [] conv_turn = ( self.discourse_manager.pure_rag_agent.generate_topic_background() ) self.conversation_history.append(conv_turn) self.knowledge_base.update_from_conv_turn( conv_turn=conv_turn, allow_create_new_node=True, insert_under_root=self.runner_argument.rag_only_baseline_mode, ) def generate_report(self) -> str: """ Generate report leveraging organized collected information in the knowledge base (i.e. mind map). The article generation follows the paradigm in STORM paper, where it considers mind map nodes as section names, and generate the report section by section. Returns: str: A string representing the report, with "#" "##" indicating hierarchical sections and [1][2] indicating references. """ with self.logging_wrapper.log_pipeline_stage( f"report generation after conv turn: {len(self.conversation_history)}" ): with self.logging_wrapper.log_event( "report generation stage: generate report" ): return self.knowledge_base.to_report() def dump_logging_and_reset(self): return self.logging_wrapper.dump_logging_and_reset() def step( self, user_utterance: str = "", simulate_user: bool = False, simulate_user_intent: str = "", ) -> ConversationTurn: """ Yields a single turn in the conversation flow. This method take a user input when user choose to inject an utterance or generates the next system utterance based on the current conversation history and defined discourse policies. It handles updating the conversation history, managing expert lists, and interacting with the knowledge base. Additionally, it logs each stage of the conversation for monitoring and debugging purposes. Args: user_utterance (str, optional): The input provided by the user. If provided, this utterance is added directly to the conversation history and returns with no further action. simulate_user (bool, optional): This is designed for automatic experiments using a LLM agent to simulate user actions. Flag indicating whether to simulate user behavior. When set to `True`, the system will generate user intents based on predefined simulation logic. Defaults to `False`. simulate_user_intent (str, optional): This is designed for automatic experiments using a LLM agent to simulate user actions. Specifies the intent to simulate for the user. This is used when `simulate_user` is `True` to guide the simulated user's responses, Returns: ConversationTurn: An object representing the latest turn in the conversation. Workflow: 1. User Utterance Handling - If `user_utterance` is provided, it is appended to the `conversation_history` 2. System Utterance Generation - If no `user_utterance` is provided, the method proceeds to generate the next system utterance. - Determines the next turn policy by consulting the `discourse_manager` with the current conversation history. - Generates a new utterance using the agent defined in the turn policy, leveraging the `knowledge_base` and `conversation_history`. - If the turn policy indicates that the experts list should be updated, it updates the expert list based on the latest utterances. 4. Knowledge Base Update - Inserts the new turn into the `knowledge_base`, optionally allowing the creation of new nodes or inserting under the root based on the `rag_only_baseline_mode` flag. - If the turn policy specifies, it reorganizes the `knowledge_base` to maintain optimal structure and relevance. """ last_conv_turn = self.conversation_history[-1] cur_turn_name = f"conv turn: {len(self.conversation_history) + 1}" with self.logging_wrapper.log_pipeline_stage( pipeline_stage=f"{cur_turn_name} stage" ): conv_turn = None if user_utterance: self.discourse_manager.next_turn_moderator_override = False conv_turn = ConversationTurn( role="Guest", raw_utterance=user_utterance, utterance_type="Original Question", ) self.conversation_history.append(conv_turn) else: with self.logging_wrapper.log_event( f"{cur_turn_name}: get turn policy" ): if self.callback_handler is not None: self.callback_handler.on_turn_policy_planning_start() turn_policy = self.discourse_manager.get_next_turn_policy( conversation_history=self.conversation_history, simulate_user=simulate_user, simulate_user_intent=simulate_user_intent, dry_run=False, ) with self.logging_wrapper.log_event( f"{cur_turn_name}: generate utterance" ): conv_turn = turn_policy.agent.generate_utterance( knowledge_base=self.knowledge_base, conversation_history=self.conversation_history, ) if turn_policy.should_update_experts_list: with self.logging_wrapper.log_event( f"{cur_turn_name}: update experts list" ): self.discourse_manager._update_expert_list_from_utterance( focus=last_conv_turn.raw_utterance, background_info=conv_turn.raw_utterance, ) if conv_turn is not None: self.conversation_history.append(conv_turn) with self.logging_wrapper.log_event( f"{cur_turn_name}: insert into knowledge base" ): if self.callback_handler is not None: self.callback_handler.on_mindmap_insert_start() self.knowledge_base.update_from_conv_turn( conv_turn=conv_turn, allow_create_new_node=True, insert_under_root=self.runner_argument.rag_only_baseline_mode, ) if self.callback_handler is not None: self.callback_handler.on_mindmap_insert_end() if turn_policy.should_reorganize_knowledge_base: with self.logging_wrapper.log_event( f"{cur_turn_name}: reorganize knowledge base" ): if self.callback_handler is not None: self.callback_handler.on_mindmap_reorg_start() self.knowledge_base.reorganize() return conv_turn ================================================ FILE: knowledge_storm/collaborative_storm/modules/__init__.py ================================================ from .article_generation import * from .grounded_question_answering import * from .grounded_question_generation import * from .information_insertion_module import * from .simulate_user import * from .warmstart_hierarchical_chat import * from .knowledge_base_summary import * from .costorm_expert_utterance_generator import * ================================================ FILE: knowledge_storm/collaborative_storm/modules/article_generation.py ================================================ import dspy from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Set, Union from .collaborative_storm_utils import clean_up_section from ...dataclass import KnowledgeBase, KnowledgeNode class ArticleGenerationModule(dspy.Module): """Use the information collected from the information-seeking conversation to write a section.""" def __init__( self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], ): super().__init__() self.write_section = dspy.Predict(WriteSection) self.engine = engine def _get_cited_information_string( self, all_citation_index: Set[int], knowledge_base: KnowledgeBase, max_words: int = 4000, ): information = [] cur_word_count = 0 for index in sorted(list(all_citation_index)): info = knowledge_base.info_uuid_to_info_dict[index] snippet = info.snippets[0] info_text = f"[{index}]: {snippet} (Question: {info.meta['question']}. Query: {info.meta['query']})" cur_snippet_length = len(info_text.split()) if cur_snippet_length + cur_word_count > max_words: break cur_word_count += cur_snippet_length information.append(info_text) return "\n".join(information) def gen_section( self, topic: str, node: KnowledgeNode, knowledge_base: KnowledgeBase ): if node is None or len(node.content) == 0: return "" if ( node.synthesize_output is not None and node.synthesize_output and not node.need_regenerate_synthesize_output ): return node.synthesize_output all_citation_index = node.collect_all_content() information = self._get_cited_information_string( all_citation_index=all_citation_index, knowledge_base=knowledge_base ) with dspy.settings.context(lm=self.engine): synthesize_output = clean_up_section( self.write_section( topic=topic, info=information, section=node.name ).output ) node.synthesize_output = synthesize_output node.need_regenerate_synthesize_output = False return node.synthesize_output def forward(self, knowledge_base: KnowledgeBase): all_nodes = knowledge_base.collect_all_nodes() node_to_paragraph = {} # Define a function to generate paragraphs for nodes def _node_generate_paragraph(node): node_gen_paragraph = self.gen_section( topic=knowledge_base.topic, node=node, knowledge_base=knowledge_base ) lines = node_gen_paragraph.split("\n") if lines[0].strip().replace("*", "").replace("#", "") == node.name: lines = lines[1:] node_gen_paragraph = "\n".join(lines) path = " -> ".join(node.get_path_from_root()) return path, node_gen_paragraph with ThreadPoolExecutor(max_workers=5) as executor: # Submit all tasks future_to_node = { executor.submit(_node_generate_paragraph, node): node for node in all_nodes } # Collect the results as they complete for future in as_completed(future_to_node): path, node_gen_paragraph = future.result() node_to_paragraph[path] = node_gen_paragraph def helper(cur_root, level): to_return = [] if cur_root is not None: hash_tag = "#" * level + " " cur_path = " -> ".join(cur_root.get_path_from_root()) node_gen_paragraph = node_to_paragraph[cur_path] to_return.append(f"{hash_tag}{cur_root.name}\n{node_gen_paragraph}") for child in cur_root.children: to_return.extend(helper(child, level + 1)) return to_return to_return = [] for child in knowledge_base.root.children: to_return.extend(helper(child, level=1)) return "\n".join(to_return) class WriteSection(dspy.Signature): """Write a Wikipedia section based on the collected information. You will be given the topic, the section you are writing and relevant information. Each information will be provided with the raw content along with question and query lead to that information. Here is the format of your writing: Use [1], [2], ..., [n] in line (for example, "The capital of the United States is Washington, D.C.[1][3]."). You DO NOT need to include a References or Sources section to list the sources at the end. """ info = dspy.InputField(prefix="The collected information:\n", format=str) topic = dspy.InputField(prefix="The topic of the page: ", format=str) section = dspy.InputField(prefix="The section you need to write: ", format=str) output = dspy.OutputField( prefix="Write the section with proper inline citations (Start your writing. Don't include the page title, section name, or try to write other sections. Do not start the section with topic name.):\n", format=str, ) ================================================ FILE: knowledge_storm/collaborative_storm/modules/callback.py ================================================ from typing import List from ...interface import Information class BaseCallbackHandler: """Base callback handler to manage callbacks from the Co-STORM pipeline.""" def on_turn_policy_planning_start(self, **kwargs): """Run when the turn policy planning begins, before deciding the direction or goal for the next conversation turn.""" pass def on_expert_action_planning_start(self, **kwargs): """Run when the expert action planning begins, preparing to determine the actions that each expert should take.""" pass def on_expert_action_planning_end(self, **kwargs): """Run when the expert action planning ends, after deciding the actions that each expert should take.""" pass def on_expert_information_collection_start(self, **kwargs): """Run when the expert information collection starts, start gathering all necessary data from selected sources.""" pass def on_expert_information_collection_end(self, info: List[Information], **kwargs): """Run when the expert information collection ends, after gathering all necessary data from selected sources.""" pass def on_expert_utterance_generation_end(self, **kwargs): """Run when the expert utterance generation ends, before creating responses or statements from each expert.""" pass def on_expert_utterance_polishing_start(self, **kwargs): """Run when the expert utterance polishing begins, to refine and improve the clarity and coherence of generated content.""" pass def on_mindmap_insert_start(self, **kwargs): """Run when the process of inserting new information into the mindmap starts.""" pass def on_mindmap_insert_end(self, **kwargs): """Run when the process of inserting new information into the mindmap ends.""" pass def on_mindmap_reorg_start(self, **kwargs): """Run when the reorganization of the mindmap begins, to restructure and optimize the flow of information.""" pass def on_expert_list_update_start(self, **kwargs): """Run when the expert list update starts, to modify or refresh the list of active experts.""" pass def on_article_generation_start(self, **kwargs): """Run when the article generation process begins, to compile and format the final article content.""" pass def on_warmstart_update(self, message, **kwargs): """Run when the warm start process has update.""" pass class LocalConsolePrintCallBackHandler(BaseCallbackHandler): def __init__(self): pass def on_turn_policy_planning_start(self, **kwargs): """Run when the turn policy planning begins, before deciding the direction or goal for the next conversation turn.""" print("Start planning next expert; inspect mind map; inspect system state.") def on_expert_action_planning_start(self, **kwargs): """Run when the expert action planning begins, preparing to determine the actions that each expert should take.""" print("Reviewing discourse history; Deciding utterance intent.") def on_expert_information_collection_start(self, **kwargs): """Run when the expert information collection ends, after gathering all necessary data from selected sources.""" print("Start searching with the search engine; browsing collected information.") def on_expert_information_collection_end(self, info: List[Information], **kwargs): """Run when the expert information collection ends, after gathering all necessary data from selected sources.""" if info: urls = [i.url for i in info] information_string = "\n".join([f"Finish browsing {url}" for url in urls]) print(information_string) def on_expert_utterance_generation_end(self, **kwargs): """Run when the expert utterance generation ends, before creating responses or statements from each expert.""" print("Finish generating utterance from collected information.") def on_expert_utterance_polishing_start(self, **kwargs): """Run when the expert utterance polishing begins, to refine and improve the clarity and coherence of generated content.""" print("Start polishing utterance.") def on_mindmap_insert_start(self, **kwargs): """Run when the process of inserting new information into the mindmap starts.""" print("Start inserting information into mind map.") def on_mindmap_insert_end(self, **kwargs): """Run when the process of inserting new information into the mindmap ends.""" print("Finish inserting information into mind map.") def on_mindmap_reorg_start(self, **kwargs): """Run when the reorganization of the mindmap begins, to restructure and optimize the flow of information.""" print("Start re-organizing mind map.") def on_expert_list_update_start(self, **kwargs): """Run when the expert list update starts, to modify or refresh the list of active experts.""" print("Start updating expert candidates.") def on_warmstart_update(self, message, **kwargs): """Run when the warm start process has update.""" print(f"Warm start update: {message}") ================================================ FILE: knowledge_storm/collaborative_storm/modules/co_storm_agents.py ================================================ import dspy from itertools import zip_longest import numpy as np from sklearn.metrics.pairwise import cosine_similarity from typing import List, Optional, TYPE_CHECKING from .callback import BaseCallbackHandler from .collaborative_storm_utils import ( extract_storm_info_snippet, _get_answer_question_module_instance, ) from .costorm_expert_utterance_generator import CoStormExpertUtteranceGenerationModule from .grounded_question_generation import GroundedQuestionGenerationModule from .simulate_user import GenSimulatedUserUtterance from ...dataclass import ConversationTurn, KnowledgeBase from ...encoder import Encoder from ...interface import Agent, Information, LMConfigs from ...logging_wrapper import LoggingWrapper if TYPE_CHECKING: from ..engine import RunnerArgument class CoStormExpert(Agent): """ Represents an expert agent in the Co-STORM framework. The `CoStormExpert` is a specialized type of `Agent` that is tasked with participating in roundtable discussions within the Co-STORM system. The expert uses language models to generate action plans, answer questions, and polish its utterances based on the current conversation history and knowledge base. It interacts with modules for action planning and question answering grounding on provided retrieval models. Args: topic (str): The conversation topic that the expert specializes in. role_name (str): The perspective of the expert's role (e.g. AI enthusiast, drug discovery expert, etc.) role_description (str): A description of the perspective of the experts lm_config (LMConfigs): Configuration for the language models runner_argument (RunnerArgument): Co-STORM runner argument logging_wrapper (LoggingWrapper): An instance of `LoggingWrapper` to log events. rm (Optional[dspy.Retrieve], optional): A retrieval module used for fetching external knowledge or context. callback_handler (BaseCallbackHandler, optional): Handles log message printing """ def __init__( self, topic: str, role_name: str, role_description: str, lm_config: LMConfigs, runner_argument: "RunnerArgument", logging_wrapper: LoggingWrapper, rm: Optional[dspy.Retrieve] = None, callback_handler: BaseCallbackHandler = None, ): super().__init__(topic, role_name, role_description) self.lm_config = lm_config self.runner_argument = runner_argument self.logging_wrapper = logging_wrapper self.callback_handler = callback_handler self.costorm_agent_utterance_generator = ( self._get_costorm_expert_utterance_generator(rm=rm) ) def _get_costorm_expert_utterance_generator( self, rm: Optional[dspy.Retrieve] = None ): return CoStormExpertUtteranceGenerationModule( action_planning_lm=self.lm_config.discourse_manage_lm, utterance_polishing_lm=self.lm_config.utterance_polishing_lm, answer_question_module=_get_answer_question_module_instance( lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=rm, ), logging_wrapper=self.logging_wrapper, callback_handler=self.callback_handler, ) def generate_utterance( self, knowledge_base: KnowledgeBase, conversation_history: List[ConversationTurn], ): with self.logging_wrapper.log_event( "CoStormExpert generate utternace: get knowledge base summary" ): if self.callback_handler is not None: self.callback_handler.on_expert_action_planning_start() conversation_summary = knowledge_base.get_knowledge_base_summary() with self.logging_wrapper.log_event( "CoStormExpert.generate_utterance generate utterance" ): last_conv_turn = conversation_history[-1] conv_turn = self.costorm_agent_utterance_generator( topic=self.topic, current_expert=self.get_role_description(), conversation_summary=conversation_summary, last_conv_turn=last_conv_turn, ).conversation_turn with self.logging_wrapper.log_event( "CoStormExpert generate utterance: polish utterance" ): if self.callback_handler is not None: self.callback_handler.on_expert_utterance_polishing_start() self.costorm_agent_utterance_generator.polish_utterance( conversation_turn=conv_turn, last_conv_turn=last_conv_turn ) return conv_turn class SimulatedUser(Agent): """ Simulated Users is a special type of Agent in Co-STORM that simulates real user interaction behavior based on the given intent. This class can be used for automatic experiments. For more information, please refer to Section 3.4 of Co-STORM paper: https://www.arxiv.org/pdf/2408.15232 """ def __init__( self, topic: str, role_name: str, role_description: str, intent: str, lm_config: LMConfigs, runner_argument: "RunnerArgument", logging_wrapper: LoggingWrapper, callback_handler: BaseCallbackHandler = None, ): super().__init__(topic, role_name, role_description) self.intent = intent self.lm_config = lm_config self.runner_argument = runner_argument self.logging_wrapper = logging_wrapper self.gen_simulated_user_utterance = GenSimulatedUserUtterance( engine=self.lm_config.question_answering_lm ) self.callback_handler = callback_handler def generate_utterance( self, knowledge_base: KnowledgeBase, conversation_history: List[ConversationTurn], ): assert ( self.intent is not None and self.intent ), "Simulate user intent is not initialized." with self.logging_wrapper.log_event( "SimulatedUser generate utternace: generate utterance" ): utterance = self.gen_simulated_user_utterance( topic=self.topic, intent=self.intent, conv_history=conversation_history ) return ConversationTurn( role="Guest", raw_utterance=utterance, utterance_type="Original Question" ) class Moderator(Agent): """ The moderator's role in the Co-STORM framework is to inject new perspectives into the conversation to avoid stagnation, repetition, or overly niche discussions. This is achieved by generating questions based on unused, uncited snippets of information retrieved since the last moderator's turn. The selected information is reranked according to its relevance to the conversation topic and its dissimilarity to the original question. The resulting top-ranked snippets are used to generate an informed question to be presented to the conversation participants. For more information, please refer to Section 3.5 of Co-STORM paper: https://www.arxiv.org/pdf/2408.15232 """ def __init__( self, topic: str, role_name: str, role_description: str, lm_config: LMConfigs, runner_argument: "RunnerArgument", logging_wrapper: LoggingWrapper, encoder: Encoder, callback_handler: BaseCallbackHandler = None, ): super().__init__(topic, role_name, role_description) self.lm_config = lm_config self.runner_argument = runner_argument self.logging_wrapper = logging_wrapper self.grounded_question_generation_module = GroundedQuestionGenerationModule( engine=self.lm_config.question_asking_lm ) self.callback_handler = callback_handler self.encoder = encoder def _get_conv_turn_unused_information( self, conv_turn: ConversationTurn, knowledge_base: KnowledgeBase ): # extract all snippets from raw retrieved information raw_retrieved_info: List[Information] = conv_turn.raw_retrieved_info raw_retrieved_single_snippet_info: List[Information] = [] for info in raw_retrieved_info: for snippet_idx in range(len(info.snippets)): raw_retrieved_single_snippet_info.append( extract_storm_info_snippet(info, snippet_index=snippet_idx) ) # get all cited information cited_info = list(knowledge_base.info_uuid_to_info_dict.values()) cited_info_hash_set = set([hash(info) for info in cited_info]) cited_snippets = [info.snippets[0] for info in cited_info] # get list of unused information unused_information: List[Information] = [ info for info in raw_retrieved_single_snippet_info if hash(info) not in cited_info_hash_set ] if not unused_information: return [] # extract snippets to get embeddings unused_information_snippets = [info.snippets[0] for info in unused_information] # get embeddings unused_snippets_embeddings = self.encoder.encode( unused_information_snippets, max_workers=20 ) claim_embedding = self.encoder.encode(conv_turn.claim_to_make) query_embedding = self.encoder.encode(conv_turn.queries) cited_snippets_embedding = self.encoder.encode(cited_snippets) # calculate similarity query_similarities = cosine_similarity( unused_snippets_embeddings, query_embedding ) max_query_similarity = np.max(query_similarities, axis=1) cited_snippets_similarity = np.max( cosine_similarity(unused_snippets_embeddings, cited_snippets_embedding), axis=1, ) cited_snippets_similarity = np.clip(cited_snippets_similarity, 0, 1) # use claim similarity to filter out "real" not useful data claim_similarity = cosine_similarity( unused_snippets_embeddings, claim_embedding.reshape(1, -1) ).flatten() claim_similarity = np.where(claim_similarity >= 0.25, 1.0, 0.0) # calculate score: snippet that is close to topic but far from query query_sim_weight = 0.5 cited_snippets_sim_weight = 1 - query_sim_weight combined_scores = ( ((1 - max_query_similarity) ** query_sim_weight) * ((1 - cited_snippets_similarity) ** cited_snippets_sim_weight) * claim_similarity ) sorted_indices = np.argsort(combined_scores)[::-1] return [unused_information[idx] for idx in sorted_indices] def _get_sorted_unused_snippets( self, knowledge_base: KnowledgeBase, conversation_history: List[ConversationTurn], last_n_conv_turn: int = 2, ): # get last N conv turn and batch encode all related strings considered_conv_turn = [] batch_snippets = [self.topic] for conv_turn in reversed(conversation_history): if len(considered_conv_turn) == last_n_conv_turn: break if conv_turn.utterance_type == "Questioning": break considered_conv_turn.append(conv_turn) batch_snippets.extend( sum([info.snippets for info in conv_turn.raw_retrieved_info], []) ) batch_snippets.append(conv_turn.claim_to_make) batch_snippets.extend(conv_turn.queries) self.encoder.encode(batch_snippets, max_workers=20) # get sorted unused snippets for each turn sorted_snippets = [] for conv_turn in considered_conv_turn: sorted_snippets.append( self._get_conv_turn_unused_information( conv_turn=conv_turn, knowledge_base=knowledge_base ) ) # use round robin rule to merge these snippets merged_snippets = [] for elements in zip_longest(*sorted_snippets, fillvalue=None): merged_snippets.extend(e for e in elements if e is not None) return merged_snippets def generate_utterance( self, knowledge_base: KnowledgeBase, conversation_history: List[ConversationTurn], ): with self.logging_wrapper.log_event( "Moderator generate utternace: get unused snippets" ): unused_snippets: List[Information] = self._get_sorted_unused_snippets( knowledge_base=knowledge_base, conversation_history=conversation_history ) with self.logging_wrapper.log_event( "Moderator generate utternace: QuestionGeneration module" ): generated_question = self.grounded_question_generation_module( topic=self.topic, knowledge_base=knowledge_base, last_conv_turn=conversation_history[-1], unused_snippets=unused_snippets, ) return ConversationTurn( role=self.role_name, raw_utterance=generated_question.raw_utterance, utterance_type="Original Question", utterance=generated_question.utterance, cited_info=generated_question.cited_info, ) class PureRAGAgent(Agent): """ PureRAGAgent only handles grounded question generation by retrieving information from the retriever based on the query. It does not utilize any other information besides the query itself. It's designed for Co-STORM paper baseline comparison. """ def __init__( self, topic: str, role_name: str, role_description: str, lm_config: LMConfigs, runner_argument: "RunnerArgument", logging_wrapper: LoggingWrapper, rm: Optional[dspy.Retrieve] = None, callback_handler: BaseCallbackHandler = None, ): super().__init__(topic, role_name, role_description) self.lm_config = lm_config self.runner_argument = runner_argument self.logging_wrapper = logging_wrapper self.grounded_question_answering_module = _get_answer_question_module_instance( lm_config=self.lm_config, runner_argument=self.runner_argument, logging_wrapper=self.logging_wrapper, rm=rm, ) def _gen_utterance_from_question(self, question: str): grounded_answer = self.grounded_question_answering_module( topic=self.topic, question=question, mode="brief", style="conversational and concise", ) conversation_turn = ConversationTurn( role=self.role_name, raw_utterance="", utterance_type="Potential Answer" ) conversation_turn.claim_to_make = question conversation_turn.raw_utterance = grounded_answer.response conversation_turn.utterance = grounded_answer.response conversation_turn.queries = grounded_answer.queries conversation_turn.raw_retrieved_info = grounded_answer.raw_retrieved_info conversation_turn.cited_info = grounded_answer.cited_info return conversation_turn def generate_topic_background(self): return self._gen_utterance_from_question(self.topic) def generate_utterance( self, knowledge_base: KnowledgeBase, conversation_history: List[ConversationTurn], ): with self.logging_wrapper.log_event( "PureRAGAgent generate utternace: generate utterance" ): return self._gen_utterance_from_question( question=conversation_history[-1].utterance ) ================================================ FILE: knowledge_storm/collaborative_storm/modules/collaborative_storm_utils.py ================================================ import dspy import os import re import sys import toml from typing import List, Tuple, Dict, Optional, TYPE_CHECKING if TYPE_CHECKING: from ..engine import RunnerArgument from ...interface import Information, Retriever, LMConfigs from ...logging_wrapper import LoggingWrapper from ...rm import BingSearch def extract_storm_info_snippet(info: Information, snippet_index: int) -> Information: """ Constructs a new Information instance with only the specified snippet index. Args: storm_info (Information): The original Information instance. snippet_index (int): The index of the snippet to retain. Returns: Information: A new Information instance with only the specified snippet. """ if snippet_index < 0 or snippet_index >= len(info.snippets): raise ValueError("Snippet index out of range") new_snippets = [info.snippets[snippet_index]] new_storm_info = Information( info.url, info.description, new_snippets, info.title, info.meta ) return new_storm_info def format_search_results( searched_results: List[Information], info_max_num_words: int = 1000, mode: str = "brief", ) -> Tuple[str, Dict[int, Information]]: """ Constructs a string from a list of search results with a specified word limit and returns a mapping of indices to Information. Args: searched_results (List[Information]): List of Information objects to process. info_max_num_words (int, optional): Maximum number of words allowed in the output string. Defaults to 1000. mode (str, optional): Mode of summarization. 'brief' takes only the first snippet of each Information. 'extensive' adds snippets iteratively until the word limit is reached. Defaults to 'brief'. Returns: Tuple[str, Dict[int, Information]]: - Formatted string with search results, constrained by the word limit. - Dictionary mapping indices to the corresponding Information objects. """ total_length = 0 extracted_snippet_queue = [] max_snippets = ( max(len(info.snippets) for info in searched_results) if searched_results else 0 ) max_snippets = 1 if mode == "brief" else max_snippets abort = False included_snippets = set() for i in range(max_snippets): for info in searched_results: if i < len(info.snippets) and not abort: cur_snippet = info.snippets[i] cur_snippet_len = len(info.snippets[i].split()) if total_length + cur_snippet_len > info_max_num_words: abort = True break if cur_snippet not in included_snippets: included_snippets.add(cur_snippet) info = extract_storm_info_snippet(info, snippet_index=i) extracted_snippet_queue.append(info) total_length += cur_snippet_len output = [] index_mapping = {} for idx, info in enumerate(extracted_snippet_queue): output.append(f"[{idx + 1}]: {info.snippets[0]}") index_mapping[idx + 1] = info assert -1 not in index_mapping return "\n".join(output), index_mapping def extract_cited_storm_info( response: str, index_to_storm_info: Dict[int, Information] ) -> Dict[int, Information]: """ Extracts a sub-dictionary of Information instances that are cited in the response. Args: response (str): The response string containing inline citations like [1], [2], etc. index_to_storm_info (Dict[int, Information]): A dictionary mapping indices to Information instances. Returns: Dict[int, Information]: A sub-dictionary with only the indices that appear in the response. """ cited_indices = set(map(int, re.findall(r"\[(\d+)\]", response))) cited_storm_info = { index: info for index, info in index_to_storm_info.items() if index in cited_indices } return cited_storm_info def trim_output_after_hint(response: str, hint: str) -> str: """ Trims the output string to only keep the substring after the given hint (not including the hint). Args: response (str): The original output string. hint (str): The hint string after which the substring should be kept. Returns: str: The trimmed output string, or the original string if the hint is not found. """ if hint in response: start_index = response.find(hint) + len(hint) return response[start_index:].strip() return response.strip("\n") def separate_citations(text: str) -> str: """ Separates multiple citations within square brackets into individual citations. Args: text (str): The input string containing citations. Returns: str: The string with separated citations. """ # Define a function to process each match def replace_citations(match): citations = match.group(1).split(",") return "".join(f"[{citation.strip()}]" for citation in citations) # Use regular expressions to find and replace citations pattern = re.compile(r"\[(\d+(?:,\s*\d+)*)\]") return pattern.sub(replace_citations, text) def extract_and_remove_citations(text: str) -> Tuple[str, List[int]]: """ Removes single inline citations from the input string and returns the modified string and a list of citation integers. Args: text (str): The input string containing citations. Returns: Tuple[str, List[int]]: The string after removal of citations and a list of citation integers. """ citations = [] # Define a function to process each match def extract_citation(match): citation = int(match.group(1)) citations.append(citation) return "" # Use regular expressions to find and replace citations pattern = re.compile(r"\[(\d+)\]") modified_text = pattern.sub(extract_citation, text) return modified_text, citations def keep_first_and_last_paragraph(text: str) -> str: """ Processes the input text to keep the first and last paragraphs and replace the middle paragraphs with '[content omitted due to space limit]'. Args: text (str): The input text containing paragraphs separated by '\n\n'. Returns: str: The processed text. """ paragraphs = text.split("\n\n") if len(paragraphs) <= 3: return text first_paragraph = paragraphs[0] last_paragraph = "\n\n".join(paragraphs[-2:]) return ( f"{first_paragraph}\n\n[content omitted due to space limit]\n\n{last_paragraph}" ) def clean_up_section(text): """Clean up a section: 1. Remove uncompleted sentences (usually due to output token limitation). 2. Deduplicate individual groups of citations. 3. Remove unnecessary summary.""" paragraphs = text.split("\n") output_paragraphs = [] summary_sec_flag = False for p in paragraphs: p = p.strip() if len(p) == 0: continue if not p.startswith("#"): p = separate_citations(p) if summary_sec_flag: if p.startswith("#"): summary_sec_flag = False else: continue if ( p.startswith("Overall") or p.startswith("In summary") or p.startswith("In conclusion") ): continue if "# Summary" in p or "# Conclusion" in p: summary_sec_flag = True continue output_paragraphs.append(p) return "\n\n".join(output_paragraphs) # Join with '\n\n' for markdown format. def load_api_key(toml_file_path): try: with open(toml_file_path, "r") as file: data = toml.load(file) except FileNotFoundError: print(f"File not found: {toml_file_path}", file=sys.stderr) return except toml.TomlDecodeError: print(f"Error decoding TOML file: {toml_file_path}", file=sys.stderr) return # Set environment variables for key, value in data.items(): os.environ[key] = str(value) def _get_answer_question_module_instance( lm_config: LMConfigs, runner_argument: "RunnerArgument", logging_wrapper: LoggingWrapper, rm: Optional[dspy.Retrieve] = None, ): from .grounded_question_answering import AnswerQuestionModule # configure retriever if rm is None: rm = BingSearch(k=runner_argument.retrieve_top_k) retriever = Retriever(rm=rm, max_thread=runner_argument.max_search_thread) # return AnswerQuestionModule instance return AnswerQuestionModule( retriever=retriever, max_search_queries=runner_argument.max_search_queries, question_answering_lm=lm_config.question_answering_lm, logging_wrapper=logging_wrapper, ) ================================================ FILE: knowledge_storm/collaborative_storm/modules/costorm_expert_utterance_generator.py ================================================ import dspy from typing import Union from .callback import BaseCallbackHandler from .collaborative_storm_utils import ( trim_output_after_hint, extract_and_remove_citations, keep_first_and_last_paragraph, ) from .grounded_question_answering import AnswerQuestionModule from .grounded_question_generation import ConvertUtteranceStyle from ...dataclass import ConversationTurn from ...logging_wrapper import LoggingWrapper class GenExpertActionPlanning(dspy.Signature): """ You are an invited speaker in the round table conversation. Your task is to make a very short note to your assistant to help you prepare for your turn in the conversation. You will be given the topic we are discussing, your expertise, and the conversation history. Take a look at conversation history, especially last few turns, then let your assistant prepare the material for you with one of following ways. 1. Original Question: Initiates a new question to other speakers. 2. Further Details: Provides additional information. 3. Information Request: Requests information from other speakers. 4. Potential Answer: Offers a possible solution or answer. Strictly follow this format: [type of contribution]: [one sentence description]. For example, Original Question: [description] """ topic = dspy.InputField(prefix="topic of discussion: ", format=str) expert = dspy.InputField(prefix="You are inivited as: ", format=str) summary = dspy.InputField(prefix="Discussion history: \n", format=str) last_utterance = dspy.InputField( prefix="Last utterance in the conversation: \n", format=str ) resposne = dspy.OutputField( prefix="Now give your note. Start with one of [Original Question, Further Details, Information Request, Potential Answer] with one sentence description\n", format=str, ) class CoStormExpertUtteranceGenerationModule(dspy.Module): def __init__( self, action_planning_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], utterance_polishing_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], answer_question_module: AnswerQuestionModule, logging_wrapper: LoggingWrapper, callback_handler: BaseCallbackHandler = None, ): self.action_planning_lm = action_planning_lm self.utterance_polishing_lm = utterance_polishing_lm self.expert_action = dspy.Predict(GenExpertActionPlanning) self.change_style = dspy.Predict(ConvertUtteranceStyle) self.answer_question_module = answer_question_module self.logging_wrapper = logging_wrapper self.callback_handler = callback_handler def parse_action(self, action): action_types = [ "Original Question", "Further Details", "Information Request", "Potential Answer", ] for action_type in action_types: if f"{action_type}:" in action: return action_type, trim_output_after_hint(action, f"{action_type}:") elif f"[{action_type}]:" in action: return action_type, trim_output_after_hint(action, f"[{action_type}]:") return "Undefined", "" def polish_utterance( self, conversation_turn: ConversationTurn, last_conv_turn: ConversationTurn ): # change utterance style action_type = conversation_turn.utterance_type with self.logging_wrapper.log_event( "RoundTableConversationModule.ConvertUtteranceStyle" ): with dspy.settings.context( lm=self.utterance_polishing_lm, show_guidelines=False ): action_string = ( f"{action_type} about: {conversation_turn.claim_to_make}" ) if action_type in ["Original Question", "Information Request"]: action_string = f"{action_type}" last_expert_utterance_wo_citation, _ = extract_and_remove_citations( last_conv_turn.utterance ) trimmed_last_expert_utterance = keep_first_and_last_paragraph( last_expert_utterance_wo_citation ) utterance = self.change_style( expert=conversation_turn.role, action=action_string, prev=trimmed_last_expert_utterance, content=conversation_turn.raw_utterance, ).utterance conversation_turn.utterance = utterance def forward( self, topic: str, current_expert: str, conversation_summary: str, last_conv_turn: ConversationTurn, ): last_utterance, _ = extract_and_remove_citations(last_conv_turn.utterance) if last_conv_turn.utterance_type in [ "Original Question", "Information Request", ]: action_type = "Potential Answer" action_content = last_utterance else: with self.logging_wrapper.log_event( "CoStormExpertUtteranceGenerationModule: GenExpertActionPlanning" ): with dspy.settings.context( lm=self.action_planning_lm, show_guidelines=False ): action = self.expert_action( topic=topic, expert=current_expert, summary=conversation_summary, last_utterance=last_utterance, ).resposne action_type, action_content = self.parse_action(action) if self.callback_handler is not None: self.callback_handler.on_expert_action_planning_end() # get response conversation_turn = ConversationTurn( role=current_expert, raw_utterance="", utterance_type=action_type ) if action_type == "Undefined": raise Exception(f"unexpected output: {action}") elif action_type in ["Further Details", "Potential Answer"]: with self.logging_wrapper.log_event( "RoundTableConversationModule: QuestionAnswering" ): grounded_answer = self.answer_question_module( topic=topic, question=action_content, mode="brief", style="conversational and concise", callback_handler=self.callback_handler, ) conversation_turn.claim_to_make = action_content conversation_turn.raw_utterance = grounded_answer.response conversation_turn.queries = grounded_answer.queries conversation_turn.raw_retrieved_info = grounded_answer.raw_retrieved_info conversation_turn.cited_info = grounded_answer.cited_info elif action_type in ["Original Question", "Information Request"]: conversation_turn.raw_utterance = action_content return dspy.Prediction(conversation_turn=conversation_turn) ================================================ FILE: knowledge_storm/collaborative_storm/modules/expert_generation.py ================================================ import dspy import re from typing import Union class GenerateExpertGeneral(dspy.Signature): """You need to select a group of diverse experts who will be suitable to be invited to a roundtable discussion on the given topic. Each expert should represent a different perspective, role, or affiliation related to this topic. You can use the background information provided about the topic for inspiration. For each expert, add a description of their expertise and what they will focus on during the discussion. No need to include speakers name in the output. Strictly follow format below: 1. [speaker 1 role]: [speaker 1 short description] 2. [speaker 2 role]: [speaker 2 short description] """ topic = dspy.InputField(prefix="Topic of interest:", format=str) background_info = dspy.InputField( prefix="Background information about the topic:\n", format=str ) topN = dspy.InputField(prefix="Number of speakers needed: ", format=str) experts = dspy.OutputField(format=str) class GenerateExpertWithFocus(dspy.Signature): """ You need to select a group of speakers who will be suitable to have roundtable discussion on the [topic] of specific [focus]. You may consider inviting speakers having opposite stands on the topic; speakers representing different interest parties; Ensure that the selected speakers are directly connected to the specific context and scenario provided. For example, if the discussion focus is about a recent event at a specific university, consider inviting students, faculty members, journalists covering the event, university officials, and local community members. Use the background information provided about the topic for inspiration. For each speaker, add a description of their interests and what they will focus on during the discussion. No need to include speakers name in the output. Strictly follow format below: 1. [speaker 1 role]: [speaker 1 short description] 2. [speaker 2 role]: [speaker 2 short description] """ topic = dspy.InputField(prefix="Topic of interest:", format=str) background_info = dspy.InputField(prefix="Background information:\n", format=str) focus = dspy.InputField(prefix="Discussion focus: ", format=str) topN = dspy.InputField(prefix="Number of speakers needed: ", format=str) experts = dspy.OutputField(format=str) class GenerateExpertModule(dspy.Module): def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.engine = engine self.generate_expert_general = dspy.Predict(GenerateExpertGeneral) self.generate_expert_w_focus = dspy.ChainOfThought(GenerateExpertWithFocus) def trim_background(self, background: str, max_words: int = 100): words = background.split() cur_len = len(words) if cur_len <= max_words: return background trimmed_words = words[: min(cur_len, max_words)] trimmed_background = " ".join(trimmed_words) return f"{trimmed_background} [rest content omitted]." def forward( self, topic: str, num_experts: int, background_info: str = "", focus: str = "" ): with dspy.settings.context(lm=self.engine, show_guidelines=False): if not focus: output = self.generate_expert_general( topic=topic, background_info=background_info, topN=num_experts ).experts else: background_info = self.trim_background( background=background_info, max_words=100 ) output = self.generate_expert_w_focus( topic=topic, background_info=background_info, focus=focus, topN=num_experts, ).experts output = output.replace("*", "").replace("[", "").replace("]", "") expert_list = [] for s in output.split("\n"): match = re.search(r"\d+\.\s*(.*)", s) if match: expert_list.append(match.group(1)) expert_list = [expert.strip() for expert in expert_list if expert.strip()] return dspy.Prediction(experts=expert_list, raw_output=output) ================================================ FILE: knowledge_storm/collaborative_storm/modules/grounded_question_answering.py ================================================ import dspy from typing import Union, List from .callback import BaseCallbackHandler from .collaborative_storm_utils import ( trim_output_after_hint, format_search_results, extract_cited_storm_info, separate_citations, ) from ...logging_wrapper import LoggingWrapper from ...utils import ArticleTextProcessing from ...interface import Information class QuestionToQuery(dspy.Signature): """You want to answer the question or support a claim using Google search. What do you type in the search box? The question is raised in a round table discussion on a topic. The question may or may not focus on the topic itself. Write the queries you will use in the following format: - query 1 - query 2 ... - query n""" topic = dspy.InputField(prefix="Topic context:", format=str) question = dspy.InputField( prefix="I want to collect information about: ", format=str ) queries = dspy.OutputField(prefix="Queries: \n", format=str) class AnswerQuestion(dspy.Signature): """You are an expert who can use information effectively. You have gathered the related information and will now use the information to form a response. Make your response as informative as possible and make sure every sentence is supported by the gathered information. If [Gathered information] is not directly related to the [Topic] and [Question], provide the most relevant answer you can based on the available information, and explain any limitations or gaps. Use [1], [2], ..., [n] in line (for example, "The capital of the United States is Washington, D.C.[1][3]."). You DO NOT need to include a References or Sources section to list the sources at the end. The style of writing should be formal. """ topic = dspy.InputField(prefix="Topic you are discussing about:", format=str) question = dspy.InputField(prefix="You want to provide insight on: ", format=str) info = dspy.InputField(prefix="Gathered information:\n", format=str) style = dspy.InputField(prefix="Style of your response should be:", format=str) answer = dspy.OutputField( prefix="Now give your response. (Try to use as many different sources as possible and do not hallucinate.)", format=str, ) class AnswerQuestionModule(dspy.Module): def __init__( self, retriever: dspy.Retrieve, max_search_queries: int, question_answering_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], logging_wrapper: LoggingWrapper, ): super().__init__() self.question_answering_lm = question_answering_lm self.question_to_query = dspy.Predict(QuestionToQuery) self.answer_question = dspy.Predict(AnswerQuestion) self.retriever = retriever self.max_search_queries = max_search_queries self.logging_wrapper = logging_wrapper def retrieve_information(self, topic, question): # decompose question to queries with self.logging_wrapper.log_event( f"AnswerQuestionModule.question_to_query ({hash(question)})" ): with dspy.settings.context(lm=self.question_answering_lm): queries = self.question_to_query(topic=topic, question=question).queries queries = trim_output_after_hint(queries, hint="Queries:") queries = [ q.replace("-", "").strip().strip('"').strip('"').strip() for q in queries.split("\n") ] queries = queries[: self.max_search_queries] self.logging_wrapper.add_query_count(count=len(queries)) with self.logging_wrapper.log_event( f"AnswerQuestionModule.retriever.retrieve ({hash(question)})" ): # retrieve information using retriever searched_results: List[Information] = self.retriever.retrieve( list(set(queries)), exclude_urls=[] ) # update storm information meta to include the question for storm_info in searched_results: storm_info.meta["question"] = question return queries, searched_results def forward( self, topic: str, question: str, mode: str = "brief", style: str = "conversational", callback_handler: BaseCallbackHandler = None, ): """ Processes a topic and question to generate a response with relevant information and citations. Args: topic (str): The topic of interest. question (str): The specific question related to the topic. mode (str, optional): Mode of summarization. 'brief' takes only the first snippet of each Information. 'extensive' adds snippets iteratively until the word limit is reached. Defaults to 'brief'. Returns: dspy.Prediction: An object containing the following: - question (str): the question to answer - queries (List[str]): List of query strings used for information retrieval. - raw_retrieved_info (List[Information]): List of Information instances retrieved. - cited_info (Dict[int, Information]): Dictionary of cited Information instances, indexed by their citation number. - response (str): The generated response string with inline citations. """ # retrieve information if callback_handler is not None: callback_handler.on_expert_information_collection_start() queries, searched_results = self.retrieve_information( topic=topic, question=question ) if callback_handler is not None: callback_handler.on_expert_information_collection_end(searched_results) # format information string for answer generation info_text, index_to_information_mapping = format_search_results( searched_results, mode=mode ) answer = "Sorry, there is insufficient information to answer the question." # generate answer to the question if info_text: with self.logging_wrapper.log_event( f"AnswerQuestionModule.answer_question ({hash(question)})" ): with dspy.settings.context( lm=self.question_answering_lm, show_guidelines=False ): answer = self.answer_question( topic=topic, question=question, info=info_text, style=style ).answer answer = ArticleTextProcessing.remove_uncompleted_sentences_with_citations( answer ) answer = trim_output_after_hint( answer, hint="Now give your response. (Try to use as many different sources as possible and do not hallucinate.)", ) # enforce single citation index bracket. [1, 2] -> [1][2] answer = separate_citations(answer) if callback_handler is not None: callback_handler.on_expert_utterance_generation_end() # construct cited search result cited_searched_results = extract_cited_storm_info( response=answer, index_to_storm_info=index_to_information_mapping ) return dspy.Prediction( question=question, queries=queries, raw_retrieved_info=searched_results, cited_info=cited_searched_results, response=answer, ) ================================================ FILE: knowledge_storm/collaborative_storm/modules/grounded_question_generation.py ================================================ """ This module handles question generation within the Co-STORM framework, specifically designed to support the Moderator role. The Moderator generates insightful, thought-provoking questions that introduce new directions into the conversation. By leveraging uncited or unused snippets of information retrieved during the discussion, the Moderator ensures the conversation remains dynamic and avoids repetitive or overly niche topics. For more detailed information, refer to Section 3.5 of the Co-STORM paper: https://www.arxiv.org/pdf/2408.15232. """ import dspy from typing import List, Union from .collaborative_storm_utils import ( format_search_results, extract_and_remove_citations, keep_first_and_last_paragraph, extract_cited_storm_info, ) from ...dataclass import ConversationTurn, KnowledgeBase from ...interface import Information class KnowledgeBaseSummmary(dspy.Signature): """Your job is to give brief summary of what's been discussed in a roundtable conversation. Contents are themantically organized into hierarchical sections. You will be presented with these sections where "#" denotes level of section. """ topic = dspy.InputField(prefix="topic: ", format=str) structure = dspy.InputField(prefix="Tree structure: \n", format=str) output = dspy.OutputField(prefix="Now give brief summary:\n", format=str) class ConvertUtteranceStyle(dspy.Signature): """ You are an invited speaker in the round table conversation. Your task is to make the question or the response more conversational and engaging to facilicate the flow of conversation. Note that this is ongoing conversation so no need to have welcoming and concluding words. Previous speaker utterance is provided only for making the conversation more natural. Note that do not hallucinate and keep the citation index like [1] as it is. Also, """ expert = dspy.InputField(prefix="You are inivited as: ", format=str) action = dspy.InputField( prefix="You want to contribute to conversation by: ", format=str ) prev = dspy.InputField(prefix="Previous speaker said: ", format=str) content = dspy.InputField( prefix="Question or response you want to say: ", format=str ) utterance = dspy.OutputField( prefix="Your utterance (keep the information as much as you can with citations, prefer shorter answers without loss of information): ", format=str, ) class GroundedQuestionGeneration(dspy.Signature): """Your job is to find next discussion focus in a roundtable conversation. You will be given previous conversation summary and some information that might assist you discover new discussion focus. Note that the new discussion focus should bring new angle and perspective to the discussion and avoid repetition. The new discussion focus should be grounded on the available information and push the boundaries of the current discussion for broader exploration. The new discussion focus should have natural flow from last utterance in the conversation. Use [1][2] in line to ground your question. """ topic = dspy.InputField(prefix="topic: ", format=str) summary = dspy.InputField(prefix="Discussion history: \n", format=str) information = dspy.InputField(prefix="Available information: \n", format=str) last_utterance = dspy.InputField( prefix="Last utterance in the conversation: \n", format=str ) output = dspy.OutputField( prefix="Now give next discussion focus in the format of one sentence question:\n", format=str, ) class GroundedQuestionGenerationModule(dspy.Module): def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.engine = engine self.gen_focus = dspy.Predict(GroundedQuestionGeneration) self.polish_style = dspy.Predict(ConvertUtteranceStyle) self.gen_summary = dspy.Predict(KnowledgeBaseSummmary) def forward( self, topic: str, knowledge_base: KnowledgeBase, last_conv_turn: ConversationTurn, unused_snippets: List[Information], ): information, index_to_information_mapping = format_search_results( unused_snippets, info_max_num_words=1000 ) summary = knowledge_base.get_knowledge_base_summary() last_utterance, _ = extract_and_remove_citations(last_conv_turn.utterance) with dspy.settings.context(lm=self.engine, show_guidelines=False): raw_utterance = self.gen_focus( topic=topic, summary=summary, information=information, last_utterance=last_utterance, ).output utterance = self.polish_style( expert="Roundtable conversation moderator", action="Raising a new question by natural transit from previous utterance.", prev=keep_first_and_last_paragraph(last_utterance), content=raw_utterance, ).utterance cited_searched_results = extract_cited_storm_info( response=utterance, index_to_storm_info=index_to_information_mapping ) return dspy.Prediction( raw_utterance=raw_utterance, utterance=utterance, cited_info=cited_searched_results, ) ================================================ FILE: knowledge_storm/collaborative_storm/modules/information_insertion_module.py ================================================ import dspy import numpy as np import re import traceback from concurrent.futures import ThreadPoolExecutor, as_completed from sklearn.metrics.pairwise import cosine_similarity from typing import List, Union, Dict, Optional from .collaborative_storm_utils import trim_output_after_hint from ...dataclass import KnowledgeNode, KnowledgeBase from ...encoder import Encoder from ...interface import Information class InsertInformation(dspy.Signature): """Your job is to insert the given information to the knowledge base. The knowledge base is a tree based data structure to organize the collection information. Each knowledge node contains information derived from themantically similar question or intent. To decide the best placement of the information, you will be navigated in this tree based data structure layer by layer. You will be presented with the question and query leads to ththeis information, and tree structure. Output should strictly follow one of options presetned below with no other information. - 'insert': to place the information under the current node. - 'step: [child node name]': to step into a specified child node. - 'create: [new child node name]': to create new child node and insert the info under it. Example outputs: - insert - step: node2 - create: node3 """ intent = dspy.InputField( prefix="Question and query leads to this info: ", format=str ) structure = dspy.InputField(prefix="Tree structure: \n", format=str) choice = dspy.OutputField(prefix="Choice:\n", format=str) class InsertInformationCandidateChoice(dspy.Signature): """Your job is to insert the given information to the knowledge base. The knowledge base is a tree based data structure to organize the collection information. Each knowledge node contains information derived from themantically similar question or intent. You will be presented with the question and query leads to this information, and candidate choices of placement. In these choices, -> denotes parent-child relationship. Note that reasonable may not be in these choices. If there exists reasonable choice, output "Best placement: [choice index]"; otherwise, output "No reasonable choice". """ intent = dspy.InputField( prefix="Question and query leads to this info: ", format=str ) choices = dspy.InputField(prefix="Candidate placement:\n", format=str) decision = dspy.OutputField(prefix="Decision:\n", format=str) class InsertInformationModule(dspy.Module): def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], encoder: Encoder): self.engine = engine self.encoder = encoder self.insert_info = dspy.ChainOfThought(InsertInformation) self.candidate_choosing = dspy.Predict(InsertInformationCandidateChoice) def _construct_intent(self, question: str, query: str): intent = "" if query == "Not applicable": return question if question: intent += f"Question: {question}\n" if query: intent += f"Query: {query}\n" if not intent: intent = "Not available." return intent def _get_navigation_choice( self, knowledge_node: KnowledgeNode, question: str, query: str ): # construct information intent intent = self._construct_intent(question, query) # construct current kb structure structure = f"Current Node: {knowledge_node.name}\n" child_names = ", ".join(knowledge_node.get_children_names()) if child_names: structure += f"Child Nodes: {child_names}" navigated_path = " -> ".join(knowledge_node.get_path_from_root()) structure += f"Path you have nagivated: {navigated_path}" # get predicted action with dspy.settings.context(lm=self.engine): predicted_action = self.insert_info( intent=intent, structure=structure ).choice # parse action cleaned_predicted_action = trim_output_after_hint( predicted_action, "Choice:" ).strip() cleaned_predicted_action = cleaned_predicted_action.strip("-").strip() if cleaned_predicted_action.startswith("insert"): return "insert", "" elif cleaned_predicted_action.startswith("step:"): node_name = trim_output_after_hint(cleaned_predicted_action, "step:") return "step", node_name elif cleaned_predicted_action.startswith("create:"): node_name = trim_output_after_hint(cleaned_predicted_action, "create:") return "create", node_name raise Exception( f"Undefined predicted action in knowledge navigation. {predicted_action}" ) def layer_by_layer_navigation_placement( self, knowledge_base: KnowledgeBase, question: str, query: str, allow_create_new_node: bool = False, root: Optional[KnowledgeNode] = None, ): current_node: KnowledgeNode = knowledge_base.root if root is None else root while True: action_type, node_name = self._get_navigation_choice( knowledge_node=current_node, question=question, query=query ) if action_type == "insert": return dspy.Prediction( information_placement=" -> ".join( current_node.get_path_from_root(root) ), note="None", ) elif action_type == "step": for child in current_node.children: if child.name == node_name: current_node = child break else: raise ValueError(f"Child node with name {node_name} not found.") elif action_type == "create": placement_path = current_node.get_path_from_root(root) if allow_create_new_node: placement_path.append(node_name) note = f"create new node: {{{node_name}}} under {{{current_node.name}}}" else: note = f"attempt to create new node: {{{node_name}}} under {{{current_node.name}}}" return dspy.Prediction( information_placement=" -> ".join(placement_path), note=note ) else: raise ValueError(f"Unknown action type: {action_type}") def _get_sorted_embed_sim_section( self, encoded_outline: np.ndarray, outlines: List[str], question: str, query: str, ): if encoded_outline is not None and encoded_outline.size > 0: encoded_query = self.encoder.encode(f"{question}, {query}") sim = cosine_similarity([encoded_query], encoded_outline)[0] sorted_indices = np.argsort(sim) sorted_outlines = np.array(outlines)[sorted_indices[::-1]] return sorted_outlines else: return outlines def _parse_selected_index(self, string: str): match = re.search(r"\[(\d+)\]", string) if match: return int(match.group(1)) try: return int(string.strip()) except: pass return None def choose_candidate_from_embedding_ranking( self, question: str, query: str, encoded_outlines: np.ndarray, outlines: List[str], top_N_candidates: int = 5, ): sorted_candidates = self._get_sorted_embed_sim_section( encoded_outlines, outlines, question, query ) considered_candidates = sorted_candidates[ : min(len(sorted_candidates), top_N_candidates) ] choices_string = "\n".join( [ f"{idx + 1}: {candidate}" for idx, candidate in enumerate(considered_candidates) ] ) with dspy.settings.context(lm=self.engine, show_guidelines=False): decision = self.candidate_choosing( intent=self._construct_intent(question=question, query=query), choices=choices_string, ).decision decision = trim_output_after_hint(decision, hint="Decision:") if "Best placement:" in decision: decision = trim_output_after_hint(decision, hint="Best placement:") selected_index = self._parse_selected_index(decision) if selected_index is not None: selected_index = selected_index - 1 if selected_index < len(sorted_candidates) and selected_index >= 0: return dspy.Prediction( information_placement=sorted_candidates[selected_index], note=f"Choosing from:\n{considered_candidates}", ) return None def _info_list_to_intent_mapping(self, information_list: List[Information]): intent_to_placement_dict = {} for info in information_list: intent = (info.meta.get("question", ""), info.meta.get("query", "")) if intent not in intent_to_placement_dict: intent_to_placement_dict[intent] = None return intent_to_placement_dict def forward( self, knowledge_base: KnowledgeBase, information: Union[Information, List[Information]], allow_create_new_node: bool = False, max_thread: int = 5, insert_root: Optional[KnowledgeNode] = None, skip_candidate_from_embedding: bool = False, ): if not isinstance(information, List): information = [information] intent_to_placement_dict: Dict = self._info_list_to_intent_mapping( information_list=information ) # process one intent def process_intent(question: str, query: str): candidate_placement = None try: if not skip_candidate_from_embedding: candidate_placement = self.choose_candidate_from_embedding_ranking( question=question, query=query, encoded_outlines=encoded_outlines, outlines=outlines, top_N_candidates=8, ) if candidate_placement is None: candidate_placement = self.layer_by_layer_navigation_placement( knowledge_base=knowledge_base, question=question, query=query, allow_create_new_node=allow_create_new_node, root=insert_root, ) return (question, query), candidate_placement except Exception as e: print(traceback.format_exc()) return (question, query), None def insert_info_to_kb(info, placement_prediction): if placement_prediction is not None: missing_node_handling = ( "raise error" if not allow_create_new_node else "create" ) knowledge_base.insert_information( path=placement_prediction.information_placement, information=info, missing_node_handling=missing_node_handling, root=insert_root, ) ( encoded_outlines, outlines, ) = knowledge_base.get_knowledge_base_structure_embedding(root=insert_root) to_return = [] if not allow_create_new_node: # use multi thread as knowledge base structure does not change with ThreadPoolExecutor(max_workers=max_thread) as executor: futures = { executor.submit(process_intent, question, query): (question, query) for (question, query) in intent_to_placement_dict } for future in as_completed(futures): (question, query), candidate_placement = future.result() intent_to_placement_dict[(question, query)] = candidate_placement # back mapping placement to each information for info in information: intent = (info.meta.get("question", ""), info.meta.get("query", "")) placement_prediction = intent_to_placement_dict.get(intent, None) insert_info_to_kb(info, placement_prediction) to_return.append((info, placement_prediction)) return to_return else: # use sequential insert as knowledge base structure might change for question, query in intent_to_placement_dict: ( encoded_outlines, outlines, ) = knowledge_base.get_knowledge_base_structure_embedding( root=insert_root ) _, placement_prediction = process_intent(question=question, query=query) intent_to_placement_dict[(question, query)] = placement_prediction for info in information: intent = (info.meta.get("question", ""), info.meta.get("query", "")) placement_prediction = intent_to_placement_dict.get(intent, None) insert_info_to_kb(info, placement_prediction) to_return.append((info, placement_prediction)) return to_return class ExpandSection(dspy.Signature): """Your task is to expand a section in the mind map by creating new subsections under the given section. You will be given a list of question and query that are used to collect information. Output should be subsection names where each section should serve as a coherent and themantic organization of information and corresponding citation numbers. These subsection names are preferred to be concise and precise. Output follows the format below: subsection 1 subsection 2 subsection 3 """ section = dspy.InputField(prefix="The section you need to expand: ", format=str) info = dspy.InputField(prefix="The collected information:\n", format=str) output = dspy.OutputField( prefix="Now provide the expanded subsection names (If there's no need to expand current section as itself serves good organization, then output None):\n", format=str, ) class ExpandNodeModule(dspy.Module): def __init__( self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], information_insert_module: dspy.Module, node_expansion_trigger_count: int, ): self.engine = engine self.expand_section = dspy.Predict(ExpandSection) self.information_insert_module = information_insert_module self.node_expansion_trigger_count = node_expansion_trigger_count def _get_cited_info_meta_string(self, node, knowledge_base): meta_string = set() for index in sorted(list(node.content)): info = knowledge_base.info_uuid_to_info_dict[index] intent = f"Question: {info.meta['question']}\nQuery: {info.meta['query']}" meta_string.add(intent) return "\n\n".join(meta_string) def _get_expand_subnode_names(self, node, knowledge_base): information = self._get_cited_info_meta_string(node, knowledge_base) node_path = node.get_path_from_root() with dspy.settings.context(lm=self.engine, show_guidelines=False): output = self.expand_section(section=node_path, info=information).output subsections = [] if "\n" in output and output != "None": subsections = output.split("\n") # remove any integer followed by a dot and a space, a leading dashline, # or a specific hint at the start of the string subsections = [ re.sub(r"^\d+\.\s|-|" + re.escape(node.name), "", text) .replace("*", "") .strip() for text in subsections ] return subsections def _find_first_node_to_expand( self, root: KnowledgeNode, expanded_nodes: List[KnowledgeNode] ): if root is None: return None if ( root not in expanded_nodes and len(root.content) >= self.node_expansion_trigger_count ): return root for child in root.children: to_return = self._find_first_node_to_expand( root=child, expanded_nodes=expanded_nodes ) if to_return is not None: return to_return return None def _expand_node(self, node: KnowledgeNode, knowledge_base: KnowledgeBase): subsection_names = self._get_expand_subnode_names(node, knowledge_base) if len(subsection_names) <= 1: return # create new nodes for subsection_name in subsection_names: # remove citation bracket in the subsection name subsection_name = re.sub(r"\[.*?\]", "", subsection_name) knowledge_base.insert_node(new_node_name=subsection_name, parent_node=node) # reset original information placement original_cited_index = node.content original_cited_information = [ knowledge_base.info_uuid_to_info_dict[index] for index in original_cited_index ] node.content = set() # re-insert under expanded section self.information_insert_module( knowledge_base=knowledge_base, information=original_cited_information, allow_create_new_node=False, insert_root=node, ) def forward(self, knowledge_base: KnowledgeBase): expanded_nodes = [] while True: node_to_expand = self._find_first_node_to_expand( root=knowledge_base.root, expanded_nodes=expanded_nodes ) if node_to_expand is None: break self._expand_node(node=node_to_expand, knowledge_base=knowledge_base) expanded_nodes.append(node_to_expand) ================================================ FILE: knowledge_storm/collaborative_storm/modules/knowledge_base_summary.py ================================================ import dspy from typing import Union from ...dataclass import KnowledgeBase class KnowledgeBaseSummmary(dspy.Signature): """Your job is to give brief summary of what's been discussed in a roundtable conversation. Contents are themantically organized into hierarchical sections. You will be presented with these sections where "#" denotes level of section. """ topic = dspy.InputField(prefix="topic: ", format=str) structure = dspy.InputField(prefix="Tree structure: \n", format=str) output = dspy.OutputField(prefix="Now give brief summary:\n", format=str) class KnowledgeBaseSummaryModule(dspy.Module): def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.engine = engine self.gen_summary = dspy.Predict(KnowledgeBaseSummmary) def forward(self, knowledge_base: KnowledgeBase): structure = knowledge_base.get_node_hierarchy_string( include_indent=False, include_full_path=False, include_hash_tag=True, include_node_content_count=False, ) with dspy.settings.context(lm=self.engine, show_guidelines=False): summary = self.gen_summary( topic=knowledge_base.topic, structure=structure ).output return summary ================================================ FILE: knowledge_storm/collaborative_storm/modules/simulate_user.py ================================================ import dspy from typing import List, Union from .collaborative_storm_utils import extract_and_remove_citations from ...dataclass import ConversationTurn from ...storm_wiki.modules.knowledge_curation import AskQuestionWithPersona class GenSimulatedUserUtterance(dspy.Module): def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.engine = engine self.ask_qeustion = dspy.Predict(AskQuestionWithPersona) def gen_conv_history_string(self, conversation_turns: List[ConversationTurn]): conv_history = [] total_turns = len(conversation_turns) for i, turn in enumerate(conversation_turns): utterance, _ = extract_and_remove_citations(turn.utterance) if i >= total_turns - 4: conv_history.append(f"{turn.role}: {utterance}") else: if turn.claim_to_make: conv_history.append(f"{turn.role}: {turn.claim_to_make}") else: conv_history.append(f"{turn.role}: {utterance}") return "\n".join(conv_history) def forward(self, topic: str, intent: str, conv_history: List[ConversationTurn]): conv_history_string = self.gen_conv_history_string(conv_history) with dspy.settings.context(lm=self.engine, show_guidelines=False): return self.ask_qeustion( topic=topic, persona=f"researcher with interest in {intent}", conv=conv_history_string, ).question ================================================ FILE: knowledge_storm/collaborative_storm/modules/warmstart_hierarchical_chat.py ================================================ """ Warm starts the Co-STORM system by conducting a background information search to establish a shared conceptual space with the user. This stage functions as a mini-STORM, where multiple LLM agents are spawned with different perspectives to engage in multi-round conversations. The knowledge base (represented as a mind map) is initialized using the information gathered during these exchanges. Additionally, the system generates a first draft of the report, which is then used to create a concise and engaging conversation. The synthesized conversation is presented to the user to help them quickly catch up on the system's current knowledge about the topic. """ import dspy import concurrent.futures from threading import Lock from typing import List, Optional, Union, TYPE_CHECKING from .callback import BaseCallbackHandler from .collaborative_storm_utils import _get_answer_question_module_instance from .expert_generation import GenerateExpertModule from .grounded_question_answering import AnswerQuestionModule from ...dataclass import ConversationTurn, KnowledgeBase from ...interface import LMConfigs from ...logging_wrapper import LoggingWrapper from ...storm_wiki.modules.outline_generation import WritePageOutline from ...utils import ArticleTextProcessing as AP if TYPE_CHECKING: from ..engine import RunnerArgument class WarmStartModerator(dspy.Signature): """ You are a moderator in a roundtable discussion. The goal is to chat with multiple experts to discuss the facts and background of the topic to familiarize the audience with the topic. You will be presented with the topic, the history of question you have already asked, and the current expert you are discussing with. Based on these information, generate the next question for the current expert to further the discussion. The output should only include the next question for the current expert. Do not include any other information or preamble. """ topic = dspy.InputField(prefix="Topic for roundtable discussion: ", format=str) history = dspy.InputField( prefix="Experts you have already interacted with: ", format=str ) current_expert = dspy.InputField(prefix="Expert you are talking with:", format=str) question = dspy.OutputField( prefix="Next question for the expert you are talking with: ", format=str ) class SectionToConvTranscript(dspy.Signature): """ You are given a section of a brief report on a specific topic. Your task is to transform this section into an engaging opening discussion for a roundtable conversation. The goal is to help participants and the audience quickly understand the key information. Both question and answer should be in the tone of roundtable discussion talking to audiences. Specifically, you need to: 1. Generate an engaging question that leverages section name and topic that opens discussion of the content. 2. Provide a brief and engaging answer (with all inline citations from original text) derived from the section serving as pointers and avoid too much details. """ topic = dspy.InputField(prefix="topic:", format=str) section_name = dspy.InputField(prefix="section name:", format=str) section_content = dspy.InputField(prefix="section content:", format=str) question = dspy.OutputField(prefix="Now give engaging question only.\nQuestion:") answer = dspy.OutputField( prefix="Now give engaging answer only with all inline citations from original text.\nAnswer:" ) class ReportToConversation(dspy.Module): def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.engine = engine self.section_to_conv_transcript = dspy.Predict(SectionToConvTranscript) def forward(self, knowledge_base: KnowledgeBase): def process_node(node, topic): with dspy.settings.context(lm=self.engine, show_guidelines=False): output = self.section_to_conv_transcript( topic=topic, section_name=node.get_path_from_root(), section_content=node.synthesize_output, ) question = output.question.replace("Question:", "").strip() answer = output.answer.replace("Answer:", "").strip() return question, answer conversations = [] nodes = knowledge_base.collect_all_nodes() nodes = [node for node in nodes if node.name != "root" and node.content] topic = knowledge_base.topic with concurrent.futures.ThreadPoolExecutor() as executor: future_to_node = { executor.submit(process_node, node, topic): node for node in nodes } for future in concurrent.futures.as_completed(future_to_node): node = future_to_node[future] question, answer = future.result() conversations.append( ConversationTurn( role="Background discussion moderator", raw_utterance=question, utterance_type="Original Question", utterance=question, cited_info=[ knowledge_base.info_uuid_to_info_dict[idx] for idx in AP.parse_citation_indices(question) ], ) ) conversations.append( ConversationTurn( role="Background discussion expert", raw_utterance=answer, utterance_type="Potential Answer", utterance=answer, cited_info=[ knowledge_base.info_uuid_to_info_dict[idx] for idx in AP.parse_citation_indices(answer) ], ) ) return conversations class WarmStartConversation(dspy.Module): def __init__( self, question_asking_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], generate_expert_module: GenerateExpertModule, answer_question_module: AnswerQuestionModule, logging_wrapper: LoggingWrapper, max_num_experts: int = 3, max_turn_per_experts: int = 2, max_thread: int = 3, callback_handler: BaseCallbackHandler = None, ): self.ask_question = dspy.Predict(WarmStartModerator) self.max_num_experts = max_num_experts self.max_turn_per_experts = max_turn_per_experts self.question_asking_lm = question_asking_lm self.answer_question_module = answer_question_module self.max_thread = max_thread self.generate_experts_module = generate_expert_module self.logging_wrapper = logging_wrapper self.callback_handler = callback_handler def format_dialogue_question_history_string( self, conversation_history: List[ConversationTurn] ): output = [] for idx, turn in enumerate(conversation_history): info = turn.claim_to_make if turn.claim_to_make else turn.utterance output.append(f"{idx + 1}: {info}") return "\n".join(output) def generate_warmstart_experts(self, topic: str): background_seeking_dialogue = self.get_background_info(topic=topic) background_info = background_seeking_dialogue.utterance gen_expert_output = self.generate_experts_module( topic=topic, background_info=background_info, num_experts=self.max_num_experts, ) return gen_expert_output.experts, background_seeking_dialogue def get_background_info(self, topic: str): question = f"Background information about {topic}" answer = self.answer_question_module( topic=topic, question=question, mode="extensive", style="conversational" ) return ConversationTurn( role="Default Background Researcher", raw_utterance=answer.response, utterance_type="Questioning", claim_to_make=question, queries=answer.queries, raw_retrieved_info=answer.raw_retrieved_info, cited_info=answer.cited_info, ) def forward(self, topic: str): with self.logging_wrapper.log_event( "warm start, perspective guided QA: identify experts" ): # do background research, generate some experts experts, background_seeking_dialogue = self.generate_warmstart_experts( topic=topic ) # init list to store the dialogue history conversation_history: List[ConversationTurn] = [] lock = Lock() # hierarchical chat: chat with one expert. Generate question, get answer def process_expert(expert): expert_name, expert_descriptoin = expert.split(":") for idx in range(self.max_turn_per_experts): with self.logging_wrapper.log_event( f"warm start, perspective guided QA: expert {expert_name}; turn {idx + 1}" ): try: with lock: history = self.format_dialogue_question_history_string( conversation_history ) with dspy.settings.context(lm=self.question_asking_lm): question = self.ask_question( topic=topic, history=history, current_expert=expert ).question answer = self.answer_question_module( topic=topic, question=question, mode="brief", style="conversational", ) conversation_turn = ConversationTurn( role=expert, claim_to_make=question, raw_utterance=answer.response, utterance_type="Support", queries=answer.queries, raw_retrieved_info=answer.raw_retrieved_info, cited_info=answer.cited_info, ) if self.callback_handler is not None: self.callback_handler.on_warmstart_update( message="\n".join( [ f"Finish browsing {url}" for url in [ i.url for i in answer.raw_retrieved_info ] ] ) ) with lock: conversation_history.append(conversation_turn) except Exception as e: print(f"Error processing expert {expert}: {e}") # multi-thread conversation with concurrent.futures.ThreadPoolExecutor( max_workers=self.max_thread ) as executor: futures = [ executor.submit(process_expert, expert) for expert in experts[: min(len(experts), self.max_num_experts)] ] concurrent.futures.wait(futures) conversation_history = [background_seeking_dialogue] + conversation_history return dspy.Prediction( conversation_history=conversation_history, experts=experts ) class GenerateWarmStartOutline(dspy.Signature): """Generate a outline of the wikipedia-like report from a roundtable discussion. You will be presented discussion points in the conversation and corresponding queries. You will be given a draft outline which you can borrow some inspiration. Do not include sections that are not mentioned in the given discussion history. Use "#" to denote section headings, "##" to denote subsection headings, and so on. Follow these guidelines: 1. Use "#" for section titles, "##" for subsection titles, "###" for subsubsection titles, and so on. 2. Do not include any additional information. 3. Exclude the topic name from the outline. The organization of outline should adopt wikiepdia style. """ topic = dspy.InputField(prefix="The topic discussed: ", format=str) draft = dspy.InputField(prefix="Draft outline you can reference to: ", format=str) conv = dspy.InputField(prefix="Discussion history:\n", format=str) outline = dspy.OutputField( prefix='Write the conversation outline (Use "#" Title" to indicate section title, "##" Title" to indicate subsection title, ...):\n', format=str, ) class GenerateWarmStartOutlineModule(dspy.Module): def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.engine = engine self.gen_outline = dspy.Predict(GenerateWarmStartOutline) self.draft_outline = dspy.Predict(WritePageOutline) def extract_questions_and_queries(self, conv: List[ConversationTurn]): context = [] for turn in conv: focus = turn.claim_to_make queries = turn.queries queries_string = "\n\t".join( f"Query {idx + 1}: {query}" for idx, query in enumerate(queries) ) string = f"Discussion focus {len(context) + 1}: {focus}\n\t{queries_string}" context.append(string) return "\n".join(context) def get_draft_outline(self, topic: str): with dspy.settings.context(lm=self.engine): return self.draft_outline(topic=topic).outline def forward(self, topic: str, conv: List[ConversationTurn]): discussion_history = self.extract_questions_and_queries(conv) draft_outline = self.get_draft_outline(topic=topic) with dspy.settings.context(lm=self.engine): outline = self.gen_outline( topic=topic, draft=draft_outline, conv=discussion_history ).outline outline = AP.clean_up_outline(outline) return dspy.Prediction(outline=outline, draft_outline=draft_outline) class WarmStartModule: def __init__( self, lm_config: LMConfigs, runner_argument: "RunnerArgument", logging_wrapper: LoggingWrapper, rm: Optional[dspy.Retrieve] = None, callback_handler: BaseCallbackHandler = None, ): generate_expert_module = GenerateExpertModule( engine=lm_config.discourse_manage_lm ) self.warmstart_conv = WarmStartConversation( question_asking_lm=lm_config.question_asking_lm, generate_expert_module=generate_expert_module, answer_question_module=_get_answer_question_module_instance( lm_config=lm_config, runner_argument=runner_argument, logging_wrapper=logging_wrapper, rm=rm, ), max_num_experts=runner_argument.warmstart_max_num_experts, max_turn_per_experts=runner_argument.warmstart_max_turn_per_experts, max_thread=runner_argument.warmstart_max_thread, logging_wrapper=logging_wrapper, callback_handler=callback_handler, ) self.warmstart_outline_gen_module = GenerateWarmStartOutlineModule( engine=lm_config.warmstart_outline_gen_lm ) self.report_to_conversation = ReportToConversation(lm_config.knowledge_base_lm) self.logging_wrapper = logging_wrapper self.callback_handler = callback_handler def initiate_warm_start(self, topic: str, knowledge_base: KnowledgeBase): """ Initiates a warm start process for the given topic by generating a warm start conversation and inserting the resulting information into a knowledge base. Args: topic (str): The topic for which to initiate the warm start process. Returns: Tuple[List[ConversationTurn], List[str], KnowledgeBase]: - A list of ConversationTurn instances representing the conversation history. - A list of strings representing the experts involved in the conversation. - A KnowledgeBase instance containing the organized information. """ warm_start_conversation_history: List[ConversationTurn] = [] warm_start_experts = None # get warm start conversations with self.logging_wrapper.log_event("warm start: perspective guided QA"): if self.callback_handler is not None: self.callback_handler.on_warmstart_update( message="Start getting familiar with the topic by chatting with multiple LLM experts (Step 1 / 4)" ) warm_start_result = self.warmstart_conv(topic=topic) warm_start_conversation_history = warm_start_result.conversation_history warm_start_experts = warm_start_result.experts # get warm start conv outline with self.logging_wrapper.log_event("warm start: outline generation"): if self.callback_handler is not None: self.callback_handler.on_warmstart_update( "Organizing collected information (Step 2 / 4)" ) warm_start_outline_output = self.warmstart_outline_gen_module( topic=topic, conv=warm_start_conversation_history ) # init knowledge base with self.logging_wrapper.log_event("warm start: insert into knowledge base"): if self.callback_handler is not None: self.callback_handler.on_warmstart_update( "Inserting collected information into knowledge base (Step 3 / 4)" ) knowledge_base.insert_from_outline_string( outline_string=warm_start_outline_output.outline ) # insert information to knowledge base for turn in warm_start_conversation_history: knowledge_base.update_from_conv_turn( conv_turn=turn, allow_create_new_node=False ) # knowledge base to report if self.callback_handler is not None: self.callback_handler.on_warmstart_update( "Synthesizing background information discussion utterances (Step 4 / 4)" ) knowledge_base.to_report() # generate engaging conversations engaging_conversations = self.report_to_conversation(knowledge_base) return ( warm_start_conversation_history, engaging_conversations, warm_start_experts, ) ================================================ FILE: knowledge_storm/dataclass.py ================================================ import dspy import numpy as np import re import threading from typing import Set, Dict, List, Optional, Union, Tuple from .encoder import Encoder from .interface import Information class ConversationTurn: """ A class to represent a turn in a conversation. Attributes: role (str): A short phrase of the role of the speaker for the current conversation turn. raw_utterance (str): The response generated by the LM model without polished style and tone. utterance_type (str): The type of utterance (e.g., statement, question). claim_to_make (Optional[str]): The point that this utterance tries to make. Should be empty if the utterance type is questioning. utterance (Optional[str]): The response generated by the model with polished style and tone. Defaults to raw_utterance if not provided. queries (List[str]): The queries used to gather information to have a grounded answer. raw_retrieved_info (List['Information']): A list of Information type that is retrieved. cited_info (Dict[int, 'Information']): A dictionary where the key is the citation index and the value is Information type. role_description (Optional[str]): A few sentences description of the role. Defaults to an empty string if not provided. """ def __init__( self, role: str, raw_utterance: str, utterance_type: str, claim_to_make: Optional[str] = None, utterance: Optional[str] = None, queries: Optional[List[str]] = None, raw_retrieved_info: Optional[List[Information]] = None, cited_info: Optional[List[Information]] = None, ): self.utterance = utterance if utterance is not None else raw_utterance self.raw_utterance = raw_utterance self.role = role if ":" not in role else role.split(":")[0] self.role_description = "" if ":" not in role else role.split(":")[1] self.queries = queries if queries is not None else [] self.raw_retrieved_info = ( raw_retrieved_info if raw_retrieved_info is not None else [] ) self.cited_info = cited_info if cited_info is not None else {} self.utterance_type = utterance_type self.claim_to_make = claim_to_make if claim_to_make is not None else "" def get_all_citation_index(self): citation_pattern = re.compile(r"\[(\d+)\]") return list(map(int, citation_pattern.findall(self.utterance))) def to_dict(self): raw_retrieved_info = [info.to_dict() for info in self.raw_retrieved_info] return { "utterance": self.utterance, "raw_utterance": self.raw_utterance, "role": self.role, "role_description": self.role_description, "queries": self.queries, "utterance_type": self.utterance_type, "claim_to_make": self.claim_to_make, "raw_retrieved_info": raw_retrieved_info, "cited_info": None, } @classmethod def from_dict(cls, conv_turn_dict: Dict): raw_retrieved_info = [ Information.from_dict(info) for info in conv_turn_dict["raw_retrieved_info"] ] return cls( utterance=conv_turn_dict["utterance"], raw_utterance=conv_turn_dict["raw_utterance"], role=f"{conv_turn_dict['role']}: {conv_turn_dict['role_description']}", queries=conv_turn_dict["queries"], raw_retrieved_info=raw_retrieved_info, cited_info=None, utterance_type=conv_turn_dict["utterance_type"], claim_to_make=conv_turn_dict["claim_to_make"], ) class KnowledgeNode: """ Class representing a node in the knowledge base. Attributes: name (str): The name of the node. content (list): A list of Information instances. children (list): A list of child KnowledgeNode instances. parent (KnowledgeNode): The parent node of the current node. """ def __init__( self, name: str, content: Optional[str] = None, parent: Optional["KnowledgeNode"] = None, children: Optional[List["KnowledgeNode"]] = None, synthesize_output: Optional[str] = None, need_regenerate_synthesize_output: bool = True, ): """ Initializes a KnowledgeNode instance. Args: name (str): The name of the node. content (list, optional): A list of information uuid. Defaults to None. parent (KnowledgeNode, optional): The parent node of the current node. Defaults to None. """ self.name = name self.content: Set[int] = set(content) if content is not None else set() self.children = [] if children is None else children self.parent = parent self.synthesize_output = synthesize_output self.need_regenerate_synthesize_output = need_regenerate_synthesize_output def collect_all_content(self): """ Collects all content from the current node and its descendants. Returns: Set[int]: A set containing all content from the current node and its descendants. """ all_content = set(self.content) for child in self.children: all_content.update(child.collect_all_content()) return all_content def has_child(self, child_node_name: str): """ Check if the node has the child of given name. """ return child_node_name in [child.name for child in self.children] def add_child(self, child_node_name: str, duplicate_handling: str = "skip"): """ Adds a child node to the current node. duplicate_handling (str): How to handle duplicate nodes. Options are "skip", "none", and "raise error". """ if self.has_child(child_node_name): if duplicate_handling == "skip": for child in self.children: if child.name == child_node_name: return child elif duplicate_handling == "raise error": raise Exception( f"Insert node error. Node {child_node_name} already exists under its parent node {self.name}." ) child_node = KnowledgeNode(name=child_node_name, parent=self) self.children.append(child_node) return child_node def get_parent(self): """ Returns the parent node of the current node. Returns: KnowledgeNode: The parent node of the current node. """ return self.parent def get_children(self): """ Returns the children of the current node. Returns: list: A list of child KnowledgeNode instances. """ return self.children def get_children_names(self): """ Returns a list of children names. """ return [child.name for child in self.children] def __repr__(self): """ Returns a string representation of the KnowledgeNode instance. Returns: str: String representation of the KnowledgeNode instance. """ return f"KnowledgeNode(name={self.name}, content={self.content}, children={len(self.children)})" def get_path_from_root(self, root: Optional["KnowledgeNode"] = None): """ Get a list of names from the root to this node. Returns: List[str]: A list of node names from the root to this node. """ path = [] current_node = self while current_node: path.append(current_node.name) if root is not None and current_node.name == root.name: break current_node = current_node.parent return path[::-1] def insert_information(self, information_index: int): if information_index not in self.content: self.need_regenerate_synthesize_output = True self.content.add(information_index) def get_all_descendents(self) -> List["KnowledgeNode"]: """ Get a list of all descendant nodes. Returns: List[KnowledgeNode]: A list of all descendant nodes. """ descendents = [] def collect_descendents(node): for child in node.children: descendents.append(child) collect_descendents(child) collect_descendents(self) return descendents def get_all_predecessors(self) -> List["KnowledgeNode"]: """ Get a list of all predecessor nodes (from current node to root). Returns: List[KnowledgeNode]: A list of all predecessor nodes. """ predecessors = [] current_node = self.parent while current_node is not None: predecessors.append(current_node) current_node = current_node.parent return predecessors def to_dict(self): """ Converts the KnowledgeNode instance to a dictionary representation. Returns: dict: The dictionary representation of the KnowledgeNode. """ return { "name": self.name, "content": list(self.content), "children": [child.to_dict() for child in self.children], "parent": self.parent.name if self.parent else None, "synthesize_output": self.synthesize_output, "need_regenerate_synthesize_output": self.need_regenerate_synthesize_output, } @classmethod def from_dict(cls, data): """ Constructs a KnowledgeNode instance from a dictionary representation. Args: data (dict): The dictionary representation of the KnowledgeNode. Returns: KnowledgeNode: The constructed KnowledgeNode instance. """ def helper(cls, data, parent_node=None): if parent_node is not None: assert data["parent"] is not None and data["parent"] == parent_node.name node = cls( name=data["name"], content=data["content"], parent=parent_node, children=None, synthesize_output=data.get("synthesize_output", None), need_regenerate_synthesize_output=data.get( "need_regenerate_synthesize_output", True ), ) for child_data in data["children"]: child_node = helper(cls, child_data, parent_node=node) node.children.append(child_node) return node return helper(cls, data) class KnowledgeBase: """ Represents the dynamic, hierarchical mind map used in Co-STORM to track and organize discourse. The knowledge base serves as a shared conceptual space between the user and the system, allowing for effective collaboration by reducing the user's cognitive load and ensuring that the discourse is easy to follow. The knowledge base is structured as a tree (or mind map) that dynamically organizes collected information and concepts as the conversation progresses. The mind map consists of concepts (nodes) and edges that represent parent-child relationships among topics. Each concept is linked to retrieved information, which is placed under the most appropriate concept based on its associated question and semantic similarity. For more details, please refer to Section 3.2 of Co-STORM paper: https://www.arxiv.org/pdf/2408.15232 Attributes: root (KnowledgeNode): The root node of the hierarchical knowledge base, representing the top-level concept. """ def __init__( self, topic: str, knowledge_base_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], node_expansion_trigger_count: int, encoder: Encoder, ): """ Initializes a KnowledgeBase instance. Args: topic (str): The topic of the knowledge base expand_node_module (dspy.Module): The module that organize knowledge base in place. The module should accept knowledge base as param. E.g. expand_node_module(self) article_generation_module (dspy.Module): The module that generate report from knowledge base. The module should return string. E.g. report = article_generation_module(self) """ from .collaborative_storm.modules.article_generation import ( ArticleGenerationModule, ) from .collaborative_storm.modules.information_insertion_module import ( InsertInformationModule, ExpandNodeModule, ) from .collaborative_storm.modules.knowledge_base_summary import ( KnowledgeBaseSummaryModule, ) self.topic: str = topic self.encoder: Encoder = encoder self.information_insert_module = InsertInformationModule( engine=knowledge_base_lm, encoder=self.encoder ) self.expand_node_module = ExpandNodeModule( engine=knowledge_base_lm, information_insert_module=self.information_insert_module, node_expansion_trigger_count=node_expansion_trigger_count, ) self.article_generation_module = ArticleGenerationModule( engine=knowledge_base_lm ) self.gen_summary_module = KnowledgeBaseSummaryModule(engine=knowledge_base_lm) self.root: KnowledgeNode = KnowledgeNode(name="root") self.kb_embedding = { "hash": hash(""), "encoded_structure": np.array([[]]), "structure_string": "", } self.info_uuid_to_info_dict: Dict[int, Information] = {} self.info_hash_to_uuid_dict: Dict[int, int] = {} self._lock = threading.Lock() def to_dict(self): info_uuid_to_info_dict = { key: value.to_dict() for key, value in self.info_uuid_to_info_dict.items() } return { "topic": self.topic, "tree": self.root.to_dict(), "info_uuid_to_info_dict": info_uuid_to_info_dict, "info_hash_to_uuid_dict": self.info_hash_to_uuid_dict, } @classmethod def from_dict( cls, data: Dict, knowledge_base_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], node_expansion_trigger_count: int, encoder: Encoder, ): knowledge_base = cls( topic=data["topic"], knowledge_base_lm=knowledge_base_lm, node_expansion_trigger_count=node_expansion_trigger_count, encoder=encoder, ) knowledge_base.root = KnowledgeNode.from_dict(data["tree"]) knowledge_base.info_hash_to_uuid_dict = { int(key): int(value) for key, value in data["info_hash_to_uuid_dict"].items() } info_uuid_to_info_dict = { int(key): Information.from_dict(value) for key, value in data["info_uuid_to_info_dict"].items() } knowledge_base.info_uuid_to_info_dict = info_uuid_to_info_dict return knowledge_base def get_knowledge_base_structure_embedding( self, root: Optional[KnowledgeNode] = None ) -> Tuple[np.ndarray, List[str]]: outline_string = self.get_node_hierarchy_string( include_indent=False, include_full_path=True, include_hash_tag=False, root=root, ) outline_string_hash = hash(outline_string) if outline_string_hash != self.kb_embedding["hash"]: outline_strings: List[str] = outline_string.split("\n") cleaned_outline_strings = [ outline.replace(" -> ", ", ") for outline in outline_strings ] encoded_outline = self.encoder.encode(cleaned_outline_strings) self.kb_embedding = { "hash": outline_string_hash, "encoded_structure": encoded_outline, "structure_string": outline_strings, } return ( self.kb_embedding["encoded_structure"], self.kb_embedding["structure_string"], ) def traverse_down(self, node): """ Traverses the tree downward from the given node. Args: node (KnowledgeNode): The node to start the traversal from. Returns: list: A list of KnowledgeNode instances in the order they were visited. """ nodes = [] def _traverse(current_node): nodes.append(current_node) for child in current_node.get_children(): _traverse(child) _traverse(node) return nodes def traverse_up(self, node): """ Traverses the tree upward from the given node. Args: node (KnowledgeNode): The node to start the traversal from. Returns: list: A list of KnowledgeNode instances in the order they were visited. """ nodes = [] while node is not None: nodes.append(node) node = node.get_parent() return nodes def collect_all_nodes(self): nodes = [] def _collect(node): nodes.append(node) for child in node.children: _collect(child) _collect(self.root) return nodes def insert_node( self, new_node_name, parent_node: Optional[KnowledgeNode] = None, duplicate_handling="skip", ): """ Inserts a new node into the knowledge base under the specified parent node. Args: new_node_name (str): The name of the new node. parent_node_name (str): The name of the parent node. If None, the new node is inserted under the root. duplicate_handling (str): How to handle duplicate nodes. Options are "skip", "none", and "raise error". """ if parent_node is None: return self.root.add_child( new_node_name, duplicate_handling=duplicate_handling ) else: return parent_node.add_child( new_node_name, duplicate_handling=duplicate_handling ) def find_node(self, current_node, node_name): """ Finds a node by name in the knowledge base. Args: current_node (KnowledgeNode): The node to start the search from. node_name (str): The name of the node to find. Returns: KnowledgeNode: The node with the specified name, or None if not found. """ if current_node.name == node_name: return current_node for child in current_node.get_children(): result = self.find_node(child, node_name) if result is not None: return result return None def insert_from_outline_string(self, outline_string, duplicate_handling="skip"): """ Creates and inserts nodes into the knowledge base from a string outline. Args: outline_string (str): The outline string where each line starts with '#' denoting the level. duplicate_handling (str): How to handle duplicate nodes. Options are "skip", "none", and "raise error". """ last_node_at_level = {} for line in outline_string.split("\n"): level = line.count("#") if level > 0: title = line.strip("# ").strip() if title.lower() in ["overview", "summary", "introduction"]: continue parent_node = None if level == 1 else last_node_at_level.get(level - 1) new_node = self.insert_node( new_node_name=title, parent_node=parent_node, duplicate_handling=duplicate_handling, ) last_node_at_level[level] = new_node for deeper_level in list(last_node_at_level.keys()): if deeper_level > level: del last_node_at_level[deeper_level] def get_node_hierarchy_string( self, include_indent=False, include_full_path=False, include_hash_tag=True, include_node_content_count=False, cited_indices: Optional[List[int]] = None, root: Optional[KnowledgeNode] = None, ) -> str: def find_node_contain_index(node, index): """ Traverses the tree downward from the given node. Args: node (KnowledgeNode): The node to start the traversal from. Returns: list: A list of KnowledgeNode instances in the order they were visited. """ nodes = [] def _traverse(current_node): if current_node is not None and index in current_node.content: nodes.append(current_node) for child in current_node.get_children(): _traverse(child) _traverse(node) return nodes paths_to_highlight = set() nodes_to_include = set() if cited_indices is not None: for index in cited_indices: for cur_node in find_node_contain_index(self.root, index): paths_to_highlight.add(" -> ".join(cur_node.get_path_from_root())) nodes_to_include.add(cur_node) nodes_to_include.update(cur_node.get_all_descendents()) predecessors = cur_node.get_all_predecessors() for predecessor in predecessors: nodes_to_include.update(predecessor.children) nodes_to_include.update(predecessors) def should_include_node(node): if cited_indices is None: return True return node in nodes_to_include def should_omit_child_nodes(node): if cited_indices is None: return False for child in node.children: if should_include_node(child): return False return True def helper(cur_root, level): to_return = [] if cur_root is not None: should_include_current_node = should_include_node(cur_root) indent = "" if not include_indent else "\t" * (level - 1) full_path = " -> ".join(cur_root.get_path_from_root(root=root)) node_info = cur_root.name if not include_full_path else full_path hash_tag = "#" * level + " " if include_hash_tag else "" content_count = ( f" ({len(cur_root.content)})" if include_node_content_count else "" ) special_note = ( "" if cited_indices is None or full_path not in paths_to_highlight else " ⭐" ) if should_include_current_node: to_return.append( f"{indent}{hash_tag}{node_info}{content_count}{special_note}" ) if should_omit_child_nodes(cur_root): if len(cur_root.children) > 0: child_indent = indent = ( "" if not include_indent else "\t" * (level) ) to_return.append(f"{child_indent}...") else: for child in cur_root.children: to_return.extend(helper(child, level + 1)) return to_return to_return = [] if root is None and self.root is not None: for child in self.root.children: to_return.extend(helper(child, level=1)) else: to_return.extend(helper(root, level=1)) return "\n".join(to_return) def find_node_by_path( self, path: str, missing_node_handling="abort", root: Optional[KnowledgeNode] = None, ): """ Returns the target node given a path string. Args: path (str): The path to the node, with node names connected by " -> ". missing_node_handling (str): How to handle missing nodes. Options are "abort", "create", and "raise error". Returns: KnowledgeNode: The target node. """ node_names = path.split(" -> ") current_node = self.root if root is None else root for name in node_names[1:]: found_node = next( (child for child in current_node.children if child.name == name), None ) if found_node is None: if missing_node_handling == "abort": return elif missing_node_handling == "create": new_node = current_node.add_child(child_node_name=name) current_node = new_node elif missing_node_handling == "raise error": structure = self.get_node_hierarchy_string( include_indent=True, include_full_path=False, include_hash_tag=True, ) raise Exception( f"Insert information error. Unable to find node {{{name}}} under {{{current_node.name}}}\n{structure}" ) else: current_node = found_node return current_node def insert_information( self, path: str, information: Information, missing_node_handling="abort", root: Optional[KnowledgeNode] = None, ): """ Inserts information into the knowledge base at the specified path. Args: path (str): The placement path string, connected by " -> " linking the name of nodes. information (Information): The information to insert. missing_node_handling (str): How to handle missing nodes. Options are "abort", "create", and "raise error". Return: uuid of insertion information """ with self._lock: target_node: KnowledgeNode = self.find_node_by_path( path=path, missing_node_handling=missing_node_handling, root=root ) information_hash = hash(information) if information.citation_uuid == -1: info_citation_uuid = self.info_hash_to_uuid_dict.get( information_hash, len(self.info_hash_to_uuid_dict) + 1 ) information.citation_uuid = info_citation_uuid self.info_hash_to_uuid_dict[information_hash] = info_citation_uuid self.info_uuid_to_info_dict[info_citation_uuid] = information if target_node is not None: self.info_uuid_to_info_dict[information.citation_uuid].meta[ "placement" ] = " -> ".join(target_node.get_path_from_root()) target_node.insert_information(information.citation_uuid) def trim_empty_leaf_nodes(self): """ Trims all leaf nodes that do not have any content. Iteratively does it until all leaf nodes have at least one content. """ def trim_node(node): if not node.children and not node.content: return True node.children = [child for child in node.children if not trim_node(child)] return not node.children and not node.content # Start the trimming process from the root while True: before_trim = len(self.get_all_leaf_nodes()) trim_node(self.root) after_trim = len(self.get_all_leaf_nodes()) if before_trim == after_trim: break def get_all_leaf_nodes(self): """ Helper function to get all leaf nodes. Returns: List[KnowledgeNode]: A list of all leaf nodes in the knowledge base. """ leaf_nodes = [] def find_leaf_nodes(node): if not node.children: leaf_nodes.append(node) for child in node.children: find_leaf_nodes(child) find_leaf_nodes(self.root) return leaf_nodes def merge_single_child_nodes(self): """ Merges content of a node with its single child and removes the child node. Iteratively does this from leaf nodes back to the root. """ def merge_node(node): # Recursively merge children first for child in node.children: merge_node(child) # If the node has exactly one child, merge its content with the child and remove the child if len(node.children) == 1: single_child = node.children[0] node.content.update(single_child.content) node.children = single_child.children for grandchild in node.children: grandchild.parent = node merge_node(self.root) def update_all_info_path(self): def _helper(node): for citation_idx in node.content: self.info_uuid_to_info_dict[citation_idx].meta["placement"] = ( " -> ".join(node.get_path_from_root()) ) for child in node.children: _helper(child) _helper(self.root) def update_from_conv_turn( self, conv_turn: ConversationTurn, allow_create_new_node: bool = False, insert_under_root: bool = False, ): if conv_turn is None: return info_to_insert = list(conv_turn.cited_info.values()) if insert_under_root: for info in info_to_insert: self.insert_information(path=self.root.name, information=info) else: self.information_insert_module( knowledge_base=self, information=info_to_insert, allow_create_new_node=allow_create_new_node, ) old_to_new_citation_idx_mapping = { old_idx: info.citation_uuid for old_idx, info in conv_turn.cited_info.items() } for old_idx, new_idx in old_to_new_citation_idx_mapping.items(): conv_turn.utterance = conv_turn.utterance.replace( f"[{old_idx}]", f"[_{new_idx}_]" ) conv_turn.raw_utterance = conv_turn.raw_utterance.replace( f"[{old_idx}]", f"[_{new_idx}_]" ) for _, new_idx in old_to_new_citation_idx_mapping.items(): conv_turn.utterance = conv_turn.utterance.replace( f"[_{new_idx}_]", f"[{new_idx}]" ) conv_turn.utterance.replace("[-1]", "") conv_turn.raw_utterance = conv_turn.raw_utterance.replace( f"[_{new_idx}_]", f"[{new_idx}]" ) conv_turn.raw_utterance.replace("[-1]", "") conv_turn.cited_info = None def get_knowledge_base_summary(self): return self.gen_summary_module(self) def reorganize(self): """ Reorganizes the knowledge base through two main processes: top-down expansion and bottom-up cleaning. The reorganization process ensures that the knowledge base remains well-structured and relevant as new information is added. It consists of the following steps: 1.Top-Down Expansion: Expands nodes that have accumulated significant amounts of information by creating subtopics, ensuring that each concept remains specific and manageable. 2.Bottom-Up Cleaning: Cleans the knowledge base by removing empty leaf nodes (nodes with no supporting information) and merging nodes that have only a single child, simplifying the structure and maintaining clarity. """ # pre-processing self.trim_empty_leaf_nodes() self.merge_single_child_nodes() # expand nodes self.expand_node_module(knowledge_base=self) # clean up self.trim_empty_leaf_nodes() self.merge_single_child_nodes() self.update_all_info_path() def to_report(self): return self.article_generation_module(knowledge_base=self) ================================================ FILE: knowledge_storm/encoder.py ================================================ import os import numpy as np from concurrent.futures import ThreadPoolExecutor, as_completed from typing import List, Tuple, Union, Optional, Dict, Literal from pathlib import Path try: import warnings with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning) if "LITELLM_LOCAL_MODEL_COST_MAP" not in os.environ: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm litellm.drop_params = True litellm.telemetry = False from litellm.caching.caching import Cache disk_cache_dir = os.path.join(Path.home(), ".storm_local_cache") litellm.cache = Cache(disk_cache_dir=disk_cache_dir, type="disk") except ImportError: class LitellmPlaceholder: def __getattr__(self, _): raise ImportError( "The LiteLLM package is not installed. Run `pip install litellm`." ) litellm = LitellmPlaceholder() class Encoder: """ A wrapper class for the LiteLLM embedding model, designed to handle embedding generation tasks efficiently. It supports parallel processing and local caching of embedding results for improved performance. The Encoder utilizes the LiteLLM library to interact with various embedding models, such as OpenAI and Azure embeddings. Users can specify the desired encoder type and provide relevant API credentials during initialization. Features: - Support for multiple embedding models (e.g., OpenAI, Azure). - Parallel processing for faster embedding generation. - Local disk caching to store and reuse embedding results. - Total token usage tracking for cost monitoring. Note: Refer to the LiteLLM documentation for details on supported embedding models: https://docs.litellm.ai/docs/embedding/supported_embedding """ def __init__( self, encoder_type: Optional[str] = None, api_key: Optional[str] = None, api_base: Optional[str] = None, api_version: Optional[str] = None, ): """ Initializes the Encoder with the appropriate embedding model. Args: encoder_type (Optional[str]): Type of encoder ('openai', 'azure', etc.). api_key (Optional[str]): API key for the encoder service. api_base (Optional[str]): API base URL for the encoder service. api_version (Optional[str]): API version for the encoder service. """ self.embedding_model_name = None self.kargs = {} self.total_token_usage = 0 # Initialize the appropriate embedding model encoder_type = encoder_type or os.getenv("ENCODER_API_TYPE") if not encoder_type: raise ValueError("ENCODER_API_TYPE environment variable is not set.") if encoder_type.lower() == "openai": self.embedding_model_name = "text-embedding-3-small" self.kargs = {"api_key": api_key or os.getenv("OPENAI_API_KEY")} elif encoder_type.lower() == "azure": self.embedding_model_name = "azure/text-embedding-3-small" self.kargs = { "api_key": api_key or os.getenv("AZURE_API_KEY"), "api_base": api_base or os.getenv("AZURE_API_BASE"), "api_version": api_version or os.getenv("AZURE_API_VERSION"), } else: raise ValueError( f"Unsupported ENCODER_API_TYPE '{encoder_type}'. Supported types are 'openai', 'azure', 'together'." ) def get_total_token_usage(self, reset: bool = False) -> int: """ Retrieves the total token usage. Args: reset (bool): If True, resets the total token usage counter after retrieval. Returns: int: The total number of tokens used. """ token_usage = self.total_token_usage if reset: self.total_token_usage = 0 return token_usage def encode(self, texts: Union[str, List[str]], max_workers: int = 5) -> np.ndarray: """ Public method to get embeddings for the given texts. Args: texts (Union[str, List[str]]): A single text string or a list of text strings to embed. Returns: np.ndarray: The array of embeddings. """ return self._get_text_embeddings(texts, max_workers=max_workers) def _get_single_text_embedding(self, text): response = litellm.embedding( model=self.embedding_model_name, input=text, caching=True, **self.kargs ) embedding = response.data[0]["embedding"] token_usage = response.get("usage", {}).get("total_tokens", 0) return text, embedding, token_usage def _get_text_embeddings( self, texts: Union[str, List[str]], max_workers: int = 5, ) -> Tuple[np.ndarray, int]: """ Get text embeddings using OpenAI's text-embedding-3-small model. Args: texts (Union[str, List[str]]): A single text string or a list of text strings to embed. max_workers (int): The maximum number of workers for parallel processing. api_key (str): The API key for accessing OpenAI's services. embedding_cache (Optional[Dict[str, np.ndarray]]): A cache to store previously computed embeddings. Returns: Tuple[np.ndarray, int]: The 2D array of embeddings and the total token usage. """ if isinstance(texts, str): _, embedding, tokens = self._get_single_text_embedding(texts) self.total_token_usage += tokens return np.array(embedding) embeddings = [] total_tokens = 0 with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = { executor.submit(self._get_single_text_embedding, text): text for text in texts } for future in as_completed(futures): try: text, embedding, tokens = future.result() embeddings.append((text, embedding, tokens)) total_tokens += tokens except Exception as e: print(f"An error occurred for text: {futures[future]}") print(e) # Sort results to match the order of the input texts embeddings.sort(key=lambda x: texts.index(x[0])) embeddings = [result[1] for result in embeddings] self.total_token_usage += total_tokens return np.array(embeddings) ================================================ FILE: knowledge_storm/interface.py ================================================ import concurrent.futures import dspy import functools import hashlib import json import logging import time from abc import ABC, abstractmethod from collections import OrderedDict from typing import Dict, List, Optional, Union, TYPE_CHECKING from .utils import ArticleTextProcessing logging.basicConfig( level=logging.INFO, format="%(name)s : %(levelname)-8s : %(message)s" ) logger = logging.getLogger(__name__) if TYPE_CHECKING: from .logging_wrapper import LoggingWrapper class InformationTable(ABC): """ The InformationTable class serves as data class to store the information collected during KnowledgeCuration stage. Create subclass to incorporate more information as needed. For example, in STORM paper https://arxiv.org/pdf/2402.14207.pdf, additional information would be perspective guided dialogue history. """ def __init__(self): pass @abstractmethod def retrieve_information(**kwargs): pass class Information: """Class to represent detailed information. Inherits from Information to include a unique identifier (URL), and extends it with a description, snippets, and title of the storm information. Attributes: description (str): Brief description. snippets (list): List of brief excerpts or snippets. title (str): The title or headline of the information. url (str): The unique URL (serving as UUID) of the information. """ def __init__(self, url, description, snippets, title, meta=None): """Initialize the Information object with detailed attributes. Args: url (str): The unique URL serving as the identifier for the information. description (str): Detailed description. snippets (list): List of brief excerpts or snippet. title (str): The title or headline of the information. """ self.description = description self.snippets = snippets self.title = title self.url = url self.meta = meta if meta is not None else {} self.citation_uuid = -1 def __hash__(self): return hash( ( self.url, tuple(sorted(self.snippets)), ) ) def __eq__(self, other): if not isinstance(other, Information): return False return ( self.url == other.url and set(self.snippets) == set(other.snippets) and self._meta_str() == other._meta_str() ) def __hash__(self): return int( self._md5_hash((self.url, tuple(sorted(self.snippets)), self._meta_str())), 16, ) def _meta_str(self): """Generate a string representation of relevant meta information.""" return f"Question: {self.meta.get('question', '')}, Query: {self.meta.get('query', '')}" def _md5_hash(self, value): """Generate an MD5 hash for a given value.""" if isinstance(value, (dict, list, tuple)): value = json.dumps(value, sort_keys=True) return hashlib.md5(str(value).encode("utf-8")).hexdigest() @classmethod def from_dict(cls, info_dict): """Create a Information object from a dictionary. Usage: info = Information.from_dict(storm_info_dict) Args: info_dict (dict): A dictionary containing keys 'url', 'description', 'snippets', and 'title' corresponding to the object's attributes. Returns: Information: An instance of Information. """ info = cls( url=info_dict["url"], description=info_dict["description"], snippets=info_dict["snippets"], title=info_dict["title"], meta=info_dict.get("meta", None), ) info.citation_uuid = int(info_dict.get("citation_uuid", -1)) return info def to_dict(self): return { "url": self.url, "description": self.description, "snippets": self.snippets, "title": self.title, "meta": self.meta, "citation_uuid": self.citation_uuid, } class ArticleSectionNode: """ The ArticleSectionNode is the dataclass for handling the section of the article. The content storage, section writing preferences are defined in this node. """ def __init__(self, section_name: str, content=None): """ section_name: section heading in string format. E.g. Introduction, History, etc. content: content of the section. Up to you for design choice of the data structure. """ self.section_name = section_name self.content = content self.children = [] self.preference = None def add_child(self, new_child_node, insert_to_front=False): if insert_to_front: self.children.insert(0, new_child_node) else: self.children.append(new_child_node) def remove_child(self, child): self.children.remove(child) class Article(ABC): def __init__(self, topic_name): self.root = ArticleSectionNode(topic_name) def find_section( self, node: ArticleSectionNode, name: str ) -> Optional[ArticleSectionNode]: """ Return the node of the section given the section name. Args: node: the node as the root to find. name: the name of node as section name Return: reference of the node or None if section name has no match """ if node.section_name == name: return node for child in node.children: result = self.find_section(child, name) if result: return result return None @abstractmethod def to_string(self) -> str: """ Export Article object into string representation. """ def get_outline_tree(self): """ Generates a hierarchical tree structure representing the outline of the document. Returns: Dict[str, Dict]: A nested dictionary representing the hierarchical structure of the document's outline. Each key is a section name, and the value is another dictionary representing the child sections, recursively forming the tree structure of the document's outline. If a section has no subsections, its value is an empty dictionary. Example: Assuming a document with a structure like: - Introduction - Background - Objective - Methods - Data Collection - Analysis The method would return: { 'Introduction': { 'Background': {}, 'Objective': {} }, 'Methods': { 'Data Collection': {}, 'Analysis': {} } } """ def build_tree(node) -> Dict[str, Dict]: tree = {} for child in node.children: tree[child.section_name] = build_tree(child) return tree if tree else {} return build_tree(self.root) def get_first_level_section_names(self) -> List[str]: """ Get first level section names """ return [i.section_name for i in self.root.children] @classmethod @abstractmethod def from_string(cls, topic_name: str, article_text: str): """ Create an instance of the Article object from a string """ pass def prune_empty_nodes(self, node=None): if node is None: node = self.root node.children[:] = [ child for child in node.children if self.prune_empty_nodes(child) ] if (node.content is None or node.content == "") and not node.children: return None else: return node class Retriever: """ An abstract base class for retriever modules. It provides a template for retrieving information based on a query. This class should be extended to implement specific retrieval functionalities. Users can design their retriever modules as needed by implementing the retrieve method. The retrieval model/search engine used for each part should be declared with a suffix '_rm' in the attribute name. """ def __init__(self, rm: dspy.Retrieve, max_thread: int = 1): self.max_thread = max_thread self.rm = rm def collect_and_reset_rm_usage(self): combined_usage = [] if hasattr(getattr(self, "rm"), "get_usage_and_reset"): combined_usage.append(getattr(self, "rm").get_usage_and_reset()) name_to_usage = {} for usage in combined_usage: for model_name, query_cnt in usage.items(): if model_name not in name_to_usage: name_to_usage[model_name] = query_cnt else: name_to_usage[model_name] += query_cnt return name_to_usage def retrieve( self, query: Union[str, List[str]], exclude_urls: List[str] = [] ) -> List[Information]: queries = query if isinstance(query, list) else [query] to_return = [] def process_query(q): retrieved_data_list = self.rm( query_or_queries=[q], exclude_urls=exclude_urls ) local_to_return = [] for data in retrieved_data_list: for i in range(len(data["snippets"])): # STORM generate the article with citations. We do not consider multi-hop citations. # Remove citations in the source to avoid confusion. data["snippets"][i] = ArticleTextProcessing.remove_citations( data["snippets"][i] ) storm_info = Information.from_dict(data) storm_info.meta["query"] = q local_to_return.append(storm_info) return local_to_return with concurrent.futures.ThreadPoolExecutor( max_workers=self.max_thread ) as executor: results = list(executor.map(process_query, queries)) for result in results: to_return.extend(result) return to_return class KnowledgeCurationModule(ABC): """ The interface for knowledge curation stage. Given topic, return collected information. """ def __init__(self, retriever: Retriever): """ Store args and finish initialization. """ self.retriever = retriever @abstractmethod def research(self, topic) -> InformationTable: """ Curate information and knowledge for the given topic Args: topic: topic of interest in natural language. Returns: collected_information: collected information in InformationTable type. """ pass class OutlineGenerationModule(ABC): """ The interface for outline generation stage. Given topic, collected information from knowledge curation stage, generate outline for the article. """ @abstractmethod def generate_outline( self, topic: str, information_table: InformationTable, **kwargs ) -> Article: """ Generate outline for the article. Required arguments include: topic: the topic of interest information_table: knowledge curation data generated from KnowledgeCurationModule More arguments could be 1. draft outline 2. user provided outline Returns: article_outline of type ArticleOutline """ pass class ArticleGenerationModule(ABC): """ The interface for article generation stage. Given topic, collected information from knowledge curation stage, generated outline from outline generation stage, """ @abstractmethod def generate_article( self, topic: str, information_table: InformationTable, article_with_outline: Article, **kwargs, ) -> Article: """ Generate article. Required arguments include: topic: the topic of interest information_table: knowledge curation data generated from KnowledgeCurationModule article_with_outline: article with specified outline from OutlineGenerationModule """ pass class ArticlePolishingModule(ABC): """ The interface for article generation stage. Given topic, collected information from knowledge curation stage, generated outline from outline generation stage, """ @abstractmethod def polish_article(self, topic: str, draft_article: Article, **kwargs) -> Article: """ Polish article. Required arguments include: topic: the topic of interest draft_article: draft article from ArticleGenerationModule. """ pass def log_execution_time(func): """Decorator to log the execution time of a function.""" @functools.wraps(func) def wrapper(self, *args, **kwargs): start_time = time.time() result = func(self, *args, **kwargs) end_time = time.time() execution_time = end_time - start_time logger.info(f"{func.__name__} executed in {execution_time:.4f} seconds") self.time[func.__name__] = execution_time return result return wrapper class LMConfigs(ABC): """Abstract base class for language model configurations of the knowledge curation engine. The language model used for each part should be declared with a suffix '_lm' in the attribute name. """ def __init__(self): pass def init_check(self): for attr_name in self.__dict__: if "_lm" in attr_name and getattr(self, attr_name) is None: logging.warning( f"Language model for {attr_name} is not initialized. Please call set_{attr_name}()" ) def collect_and_reset_lm_history(self): history = [] for attr_name in self.__dict__: if "_lm" in attr_name and hasattr(getattr(self, attr_name), "history"): history.extend(getattr(self, attr_name).history) getattr(self, attr_name).history = [] return history def collect_and_reset_lm_usage(self): combined_usage = [] for attr_name in self.__dict__: if "_lm" in attr_name and hasattr( getattr(self, attr_name), "get_usage_and_reset" ): combined_usage.append(getattr(self, attr_name).get_usage_and_reset()) model_name_to_usage = {} for usage in combined_usage: for model_name, tokens in usage.items(): if model_name not in model_name_to_usage: model_name_to_usage[model_name] = tokens else: model_name_to_usage[model_name]["prompt_tokens"] += tokens[ "prompt_tokens" ] model_name_to_usage[model_name]["completion_tokens"] += tokens[ "completion_tokens" ] return model_name_to_usage def log(self): return OrderedDict( { attr_name: getattr(self, attr_name).kwargs for attr_name in self.__dict__ if "_lm" in attr_name and hasattr(getattr(self, attr_name), "kwargs") } ) class Engine(ABC): def __init__(self, lm_configs: LMConfigs): self.lm_configs = lm_configs self.time = {} self.lm_cost = {} # Cost of language models measured by in/out tokens. self.rm_cost = {} # Cost of retrievers measured by number of queries. def log_execution_time_and_lm_rm_usage(self, func): """Decorator to log the execution time, language model usage, and retrieval model usage of a function.""" @functools.wraps(func) def wrapper(*args, **kwargs): start_time = time.time() result = func(*args, **kwargs) end_time = time.time() execution_time = end_time - start_time self.time[func.__name__] = execution_time logger.info(f"{func.__name__} executed in {execution_time:.4f} seconds") self.lm_cost[func.__name__] = self.lm_configs.collect_and_reset_lm_usage() if hasattr(self, "retriever"): self.rm_cost[func.__name__] = ( self.retriever.collect_and_reset_rm_usage() ) return result return wrapper def apply_decorators(self): """Apply decorators to methods that need them.""" methods_to_decorate = [ method_name for method_name in dir(self) if callable(getattr(self, method_name)) and method_name.startswith("run_") ] for method_name in methods_to_decorate: original_method = getattr(self, method_name) decorated_method = self.log_execution_time_and_lm_rm_usage(original_method) setattr(self, method_name, decorated_method) @abstractmethod def run_knowledge_curation_module(self, **kwargs) -> Optional[InformationTable]: pass @abstractmethod def run_outline_generation_module(self, **kwarg) -> Article: pass @abstractmethod def run_article_generation_module(self, **kwarg) -> Article: pass @abstractmethod def run_article_polishing_module(self, **kwarg) -> Article: pass @abstractmethod def run(self, **kwargs): pass def summary(self): print("***** Execution time *****") for k, v in self.time.items(): print(f"{k}: {v:.4f} seconds") print("***** Token usage of language models: *****") for k, v in self.lm_cost.items(): print(f"{k}") for model_name, tokens in v.items(): print(f" {model_name}: {tokens}") print("***** Number of queries of retrieval models: *****") for k, v in self.rm_cost.items(): print(f"{k}: {v}") def reset(self): self.time = {} self.lm_cost = {} self.rm_cost = {} class Agent(ABC): """ Interface for STORM and Co-STORM LLM agent This class must be implemented by any subclass of `Agent` to define how the agent generates an utterance. The generated utterance can be influenced by the conversation history, knowledge base, and any additional parameters passed via `kwargs`. The implementation should align with the specific role and perspective of the agent, as defined by the agent's topic, role name, and role description. Args: knowledge_base (KnowledgeBase): The current knowledge base (e.g., mind map in Co-STORM) that contains the accumulated information relevant to the conversation. conversation_history (List[ConversationTurn]): A list of past conversation turns, providing context for generating the next utterance. The agent can refer to this history to maintain continuity and relevance in the conversation. logging_wrapper (LoggingWrapper): A wrapper used for logging important events during the utterance generation process. **kwargs: Additional arguments that can be passed to the method for more specialized utterance generation behavior depending on the agent's specific implementation. Returns: ConversationTurn: A new conversation turn generated by the agent, containing the agent's response, including the role, utterance type, and relevant information from the knowledge base. Notes: - Subclasses of `Agent` should define the exact strategy for generating the utterance, which could involve interacting with a language model, retrieving relevant knowledge, or following specific conversational policies. - The agent's role, perspective, and the knowledge base content will influence how the utterance is formulated. """ from .dataclass import KnowledgeBase, ConversationTurn def __init__(self, topic: str, role_name: str, role_description: str): self.topic = topic self.role_name = role_name self.role_description = role_description def get_role_description(self): if self.role_description: return f"{self.role_name}: {self.role_description}" return self.role_name @abstractmethod def generate_utterance( self, knowledge_base: KnowledgeBase, conversation_history: List[ConversationTurn], logging_wrapper: "LoggingWrapper", **kwargs, ): pass ================================================ FILE: knowledge_storm/lm.py ================================================ import backoff import dspy import functools import logging import os import random import requests import threading from typing import Optional, Literal, Any import ujson from pathlib import Path from dsp import ERRORS, backoff_hdlr, giveup_hdlr from dsp.modules.hf import openai_to_hf from dsp.modules.hf_client import send_hftgi_request_v01_wrapped from openai import OpenAI, AzureOpenAI from transformers import AutoTokenizer try: from anthropic import RateLimitError except ImportError: RateLimitError = None ############################ # Code copied from https://github.com/stanfordnlp/dspy/blob/main/dspy/clients/lm.py on Sep 29, 2024 # try: import warnings with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning) if "LITELLM_LOCAL_MODEL_COST_MAP" not in os.environ: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm litellm.drop_params = True litellm.telemetry = False from litellm.caching.caching import Cache disk_cache_dir = os.path.join(Path.home(), ".storm_local_cache") litellm.cache = Cache(disk_cache_dir=disk_cache_dir, type="disk") # except ImportError: # class LitellmPlaceholder: # def __getattr__(self, _): # raise ImportError( # "The LiteLLM package is not installed. Run `pip install litellm`." # ) # litellm = LitellmPlaceholder() LM_LRU_CACHE_MAX_SIZE = 3000 class LM: def __init__( self, model, model_type="chat", temperature=0.0, max_tokens=1000, cache=True, **kwargs, ): self.model = model self.model_type = model_type self.cache = cache self.kwargs = dict(temperature=temperature, max_tokens=max_tokens, **kwargs) self.history = [] if "o1-" in model: assert ( max_tokens >= 5000 and temperature == 1.0 ), "OpenAI's o1-* models require passing temperature=1.0 and max_tokens >= 5000 to `dspy.LM(...)`" def __call__(self, prompt=None, messages=None, **kwargs): # Build the request. cache = kwargs.pop("cache", self.cache) messages = messages or [{"role": "user", "content": prompt}] kwargs = {**self.kwargs, **kwargs} # Make the request and handle LRU & disk caching. if self.model_type == "chat": completion = cached_litellm_completion if cache else litellm_completion else: completion = ( cached_litellm_text_completion if cache else litellm_text_completion ) response = completion( ujson.dumps(dict(model=self.model, messages=messages, **kwargs)) ) outputs = [ c.message.content if hasattr(c, "message") else c["text"] for c in response["choices"] ] # Logging, with removed api key & where `cost` is None on cache hit. kwargs = {k: v for k, v in kwargs.items() if not k.startswith("api_")} entry = dict(prompt=prompt, messages=messages, kwargs=kwargs, response=response) entry = dict(**entry, outputs=outputs, usage=dict(response["usage"])) entry = dict( **entry, cost=response.get("_hidden_params", {}).get("response_cost") ) self.history.append(entry) return outputs def inspect_history(self, n: int = 1): _inspect_history(self, n) @functools.lru_cache(maxsize=LM_LRU_CACHE_MAX_SIZE) def cached_litellm_completion(request): return litellm_completion(request, cache={"no-cache": False, "no-store": False}) def litellm_completion(request, cache={"no-cache": True, "no-store": True}): kwargs = ujson.loads(request) return litellm.completion(cache=cache, **kwargs) @functools.lru_cache(maxsize=LM_LRU_CACHE_MAX_SIZE) def cached_litellm_text_completion(request): return litellm_text_completion( request, cache={"no-cache": False, "no-store": False} ) def litellm_text_completion(request, cache={"no-cache": True, "no-store": True}): kwargs = ujson.loads(request) # Extract the provider and model from the model string. model = kwargs.pop("model").split("/", 1) provider, model = model[0] if len(model) > 1 else "openai", model[-1] # Use the API key and base from the kwargs, or from the environment. api_key = kwargs.pop("api_key", None) or os.getenv(f"{provider}_API_KEY") api_base = kwargs.pop("api_base", None) or os.getenv(f"{provider}_API_BASE") # Build the prompt from the messages. prompt = "\n\n".join( [x["content"] for x in kwargs.pop("messages")] + ["BEGIN RESPONSE:"] ) return litellm.text_completion( cache=cache, model=f"text-completion-openai/{model}", api_key=api_key, api_base=api_base, prompt=prompt, **kwargs, ) def _green(text: str, end: str = "\n"): return "\x1b[32m" + str(text).lstrip() + "\x1b[0m" + end def _red(text: str, end: str = "\n"): return "\x1b[31m" + str(text) + "\x1b[0m" + end def _inspect_history(lm, n: int = 1): """Prints the last n prompts and their completions.""" for item in lm.history[-n:]: messages = item["messages"] or [{"role": "user", "content": item["prompt"]}] outputs = item["outputs"] print("\n\n\n") for msg in messages: print(_red(f"{msg['role'].capitalize()} message:")) print(msg["content"].strip()) print("\n") print(_red("Response:")) print(_green(outputs[0].strip())) if len(outputs) > 1: choices_text = f" \t (and {len(outputs)-1} other completions)" print(_red(choices_text, end="")) print("\n\n\n") ############################ class LitellmModel(LM): """A wrapper class for LiteLLM. Check out https://docs.litellm.ai/docs/providers for usage details. """ def __init__( self, model: str = "openai/gpt-4o-mini", api_key: Optional[str] = None, model_type: Literal["chat", "text"] = "chat", **kwargs, ): super().__init__(model=model, api_key=api_key, model_type=model_type, **kwargs) self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 def log_usage(self, response): """Log the total tokens from the OpenAI API response.""" usage_data = response.get("usage") if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.get("prompt_tokens", 0) self.completion_tokens += usage_data.get("completion_tokens", 0) def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.model or self.kwargs.get("model") or self.kwargs.get("engine"): { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage def __call__(self, prompt=None, messages=None, **kwargs): # Build the request. cache = kwargs.pop("cache", self.cache) messages = messages or [{"role": "user", "content": prompt}] kwargs = {**self.kwargs, **kwargs} # Make the request and handle LRU & disk caching. if self.model_type == "chat": completion = cached_litellm_completion if cache else litellm_completion else: completion = ( cached_litellm_text_completion if cache else litellm_text_completion ) response = completion( ujson.dumps(dict(model=self.model, messages=messages, **kwargs)) ) response_dict = response.json() self.log_usage(response_dict) outputs = [ c.message.content if hasattr(c, "message") else c["text"] for c in response["choices"] ] # Logging, with removed api key & where `cost` is None on cache hit. kwargs = {k: v for k, v in kwargs.items() if not k.startswith("api_")} entry = dict( prompt=prompt, messages=messages, kwargs=kwargs, response=response_dict ) entry = dict(**entry, outputs=outputs, usage=dict(response_dict["usage"])) entry = dict( **entry, cost=response.get("_hidden_params", {}).get("response_cost") ) self.history.append(entry) return outputs # ======================================================================== # The following language model classes were deprecated after v1.1.0. # They remain in this file for backward compatibility but will no longer be maintained. class OpenAIModel(dspy.OpenAI): """A wrapper class for dspy.OpenAI.""" def __init__( self, model: str = "gpt-4o-mini", api_key: Optional[str] = None, model_type: Literal["chat", "text"] = None, **kwargs, ): super().__init__(model=model, api_key=api_key, model_type=model_type, **kwargs) self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 def log_usage(self, response): """Log the total tokens from the OpenAI API response.""" usage_data = response.get("usage") if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.get("prompt_tokens", 0) self.completion_tokens += usage_data.get("completion_tokens", 0) def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.kwargs.get("model") or self.kwargs.get("engine"): { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage def __call__( self, prompt: str, only_completed: bool = True, return_sorted: bool = False, **kwargs, ) -> list[dict[str, Any]]: """Copied from dspy/dsp/modules/gpt3.py with the addition of tracking token usage.""" assert only_completed, "for now" assert return_sorted is False, "for now" # if kwargs.get("n", 1) > 1: # if self.model_type == "chat": # kwargs = {**kwargs} # else: # kwargs = {**kwargs, "logprobs": 5} response = self.request(prompt, **kwargs) # Log the token usage from the OpenAI API response. self.log_usage(response) choices = response["choices"] completed_choices = [c for c in choices if c["finish_reason"] != "length"] if only_completed and len(completed_choices): choices = completed_choices completions = [self._get_choice_text(c) for c in choices] if return_sorted and kwargs.get("n", 1) > 1: scored_completions = [] for c in choices: tokens, logprobs = ( c["logprobs"]["tokens"], c["logprobs"]["token_logprobs"], ) if "<|endoftext|>" in tokens: index = tokens.index("<|endoftext|>") + 1 tokens, logprobs = tokens[:index], logprobs[:index] avglog = sum(logprobs) / len(logprobs) scored_completions.append((avglog, self._get_choice_text(c))) scored_completions = sorted(scored_completions, reverse=True) completions = [c for _, c in scored_completions] return completions class DeepSeekModel(dspy.OpenAI): """A wrapper class for DeepSeek API, compatible with dspy.OpenAI.""" def __init__( self, model: str = "deepseek-chat", api_key: Optional[str] = None, api_base: str = "https://api.deepseek.com", **kwargs, ): super().__init__(model=model, api_key=api_key, api_base=api_base, **kwargs) self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 self.model = model self.api_key = api_key or os.getenv("DEEPSEEK_API_KEY") self.api_base = api_base if not self.api_key: raise ValueError( "DeepSeek API key must be provided either as an argument or as an environment variable DEEPSEEK_API_KEY" ) def log_usage(self, response): """Log the total tokens from the DeepSeek API response.""" usage_data = response.get("usage") if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.get("prompt_tokens", 0) self.completion_tokens += usage_data.get("completion_tokens", 0) def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.model: { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage @backoff.on_exception( backoff.expo, ERRORS, max_time=1000, on_backoff=backoff_hdlr, giveup=giveup_hdlr, ) def _create_completion(self, prompt: str, **kwargs): """Create a completion using the DeepSeek API.""" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}", } data = { "model": self.model, "messages": [{"role": "user", "content": prompt}], **kwargs, } response = requests.post( f"{self.api_base}/v1/chat/completions", headers=headers, json=data ) response.raise_for_status() return response.json() def __call__( self, prompt: str, only_completed: bool = True, return_sorted: bool = False, **kwargs, ) -> list[dict[str, Any]]: """Call the DeepSeek API to generate completions.""" assert only_completed, "for now" assert return_sorted is False, "for now" response = self._create_completion(prompt, **kwargs) # Log the token usage from the DeepSeek API response. self.log_usage(response) choices = response["choices"] completions = [choice["message"]["content"] for choice in choices] history = { "prompt": prompt, "response": response, "kwargs": kwargs, } self.history.append(history) return completions class AzureOpenAIModel(dspy.LM): """A wrapper class of Azure OpenAI endpoint. Note: param::model should match the deployment_id on your Azure platform. """ def __init__( self, azure_endpoint: str, api_version: str, model: str, api_key: str, model_type: Literal["chat", "text"] = "chat", **kwargs, ): super().__init__(model=model) self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 self.model = model self.provider = "azure" self.model_type = model_type self.client = AzureOpenAI( azure_endpoint=azure_endpoint, api_key=api_key, api_version=api_version, ) self.prompt_tokens = 0 self.completion_tokens = 0 self.kwargs = { "model": model, "temperature": 0.0, "max_tokens": 150, "top_p": 1, "frequency_penalty": 0, "presence_penalty": 0, "n": 1, **kwargs, } @backoff.on_exception( backoff.expo, ERRORS, max_time=1000, on_backoff=backoff_hdlr, giveup=giveup_hdlr, ) def basic_request(self, prompt: str, **kwargs) -> Any: kwargs = {**self.kwargs, **kwargs} try: if self.model_type == "chat": messages = [{"role": "user", "content": prompt}] response = self.client.chat.completions.create( messages=messages, **kwargs ) else: response = self.client.completions.create(prompt=prompt, **kwargs) self.log_usage(response) history_entry = { "prompt": prompt, "response": dict(response), "kwargs": kwargs, } self.history.append(history_entry) return response except Exception as e: logging.error(f"Error making request to Azure OpenAI: {str(e)}") raise def _get_choice_text(self, choice: Any) -> str: """Extract text from a choice object based on model type.""" if self.model_type == "chat": return choice.message.content return choice.text def log_usage(self, response): """Log the total tokens from response.""" usage_data = response.usage if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.prompt_tokens self.completion_tokens += usage_data.completion_tokens def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.model: { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage def __call__( self, prompt: str, only_completed: bool = True, return_sorted: bool = False, **kwargs, ) -> list[str]: """Get completions from Azure OpenAI. Args: prompt: The prompt to send to the model only_completed: Only return completed responses return_sorted: Sort completions by probability (not implemented) **kwargs: Additional arguments to pass to the API Returns: List of completion strings """ response = self.basic_request(prompt, **kwargs) choices = response.choices completed_choices = [c for c in choices if c.finish_reason != "length"] if only_completed and completed_choices: choices = completed_choices completions = [self._get_choice_text(c) for c in choices] return completions class GroqModel(dspy.OpenAI): """A wrapper class for Groq API (https://console.groq.com/), compatible with dspy.OpenAI.""" def __init__( self, model: str = "llama3-70b-8192", api_key: Optional[str] = None, api_base: str = "https://api.groq.com/openai/v1", **kwargs, ): super().__init__(model=model, api_key=api_key, api_base=api_base, **kwargs) self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 self.model = model self.api_key = api_key or os.getenv("GROQ_API_KEY") self.api_base = api_base if not self.api_key: raise ValueError( "Groq API key must be provided either as an argument or as an environment variable GROQ_API_KEY" ) def log_usage(self, response): """Log the total tokens from the Groq API response.""" usage_data = response.get("usage") if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.get("prompt_tokens", 0) self.completion_tokens += usage_data.get("completion_tokens", 0) def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.model: { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage @backoff.on_exception( backoff.expo, ERRORS, max_time=1000, on_backoff=backoff_hdlr, giveup=giveup_hdlr, ) def _create_completion(self, prompt: str, **kwargs): """Create a completion using the Groq API.""" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}", } # Remove unsupported fields kwargs.pop("logprobs", None) kwargs.pop("logit_bias", None) kwargs.pop("top_logprobs", None) # Ensure N is 1 if supplied if "n" in kwargs and kwargs["n"] != 1: raise ValueError("Groq API only supports N=1") # Adjust temperature if it's 0 if kwargs.get("temperature", 1) == 0: kwargs["temperature"] = 1e-8 data = { "model": self.model, "messages": [{"role": "user", "content": prompt}], **kwargs, } # Remove 'name' field from messages if present for message in data["messages"]: message.pop("name", None) response = requests.post( f"{self.api_base}/chat/completions", headers=headers, json=data ) response.raise_for_status() return response.json() def __call__( self, prompt: str, only_completed: bool = True, return_sorted: bool = False, **kwargs, ) -> list[dict[str, Any]]: """Call the Groq API to generate completions.""" assert only_completed, "for now" assert return_sorted is False, "for now" response = self._create_completion(prompt, **kwargs) # Log the token usage from the Groq API response. self.log_usage(response) choices = response["choices"] completions = [choice["message"]["content"] for choice in choices] history = { "prompt": prompt, "response": response, "kwargs": kwargs, } self.history.append(history) return completions class ClaudeModel(dspy.dsp.modules.lm.LM): """Copied from dspy/dsp/modules/anthropic.py with the addition of tracking token usage.""" def __init__( self, model: str, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs, ): super().__init__(model) try: from anthropic import Anthropic except ImportError as err: raise ImportError("Claude requires `pip install anthropic`.") from err self.provider = "anthropic" self.api_key = api_key = ( os.environ.get("ANTHROPIC_API_KEY") if api_key is None else api_key ) self.api_base = ( "https://api.anthropic.com/v1/messages" if api_base is None else api_base ) self.kwargs = { "temperature": kwargs.get("temperature", 0.0), "max_tokens": min(kwargs.get("max_tokens", 4096), 4096), "top_p": kwargs.get("top_p", 1.0), "top_k": kwargs.get("top_k", 1), "n": kwargs.pop("n", kwargs.pop("num_generations", 1)), **kwargs, "model": model, } self.history: list[dict[str, Any]] = [] self.client = Anthropic(api_key=api_key) self.model = model self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 def log_usage(self, response): """Log the total tokens from the Anthropic API response.""" usage_data = response.usage if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.input_tokens self.completion_tokens += usage_data.output_tokens def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.model: { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage def basic_request(self, prompt: str, **kwargs): raw_kwargs = kwargs kwargs = {**self.kwargs, **kwargs} # caching mechanism requires hashable kwargs kwargs["messages"] = [{"role": "user", "content": prompt}] kwargs.pop("n") response = self.client.messages.create(**kwargs) # history = { # "prompt": prompt, # "response": response, # "kwargs": kwargs, # "raw_kwargs": raw_kwargs, # } json_serializable_history = { "prompt": prompt, "response": { "content": response.content[0].text, "model": response.model, "role": response.role, "stop_reason": response.stop_reason, "stop_sequence": response.stop_sequence, "type": response.type, "usage": { "input_tokens": response.usage.input_tokens, "output_tokens": response.usage.output_tokens, }, }, "kwargs": kwargs, "raw_kwargs": raw_kwargs, } self.history.append(json_serializable_history) return response @backoff.on_exception( backoff.expo, (RateLimitError,), max_time=1000, max_tries=8, on_backoff=backoff_hdlr, giveup=giveup_hdlr, ) def request(self, prompt: str, **kwargs): """Handles retrieval of completions from Anthropic whilst handling API errors.""" return self.basic_request(prompt, **kwargs) def __call__(self, prompt, only_completed=True, return_sorted=False, **kwargs): """Retrieves completions from Anthropic. Args: prompt (str): prompt to send to Anthropic only_completed (bool, optional): return only completed responses and ignores completion due to length. Defaults to True. return_sorted (bool, optional): sort the completion choices using the returned probabilities. Defaults to False. Returns: list[str]: list of completion choices """ assert only_completed, "for now" assert return_sorted is False, "for now" # per eg here: https://docs.anthropic.com/claude/reference/messages-examples # max tokens can be used as a proxy to return smaller responses # so this cannot be a proper indicator for incomplete response unless it isnt the user-intent. n = kwargs.pop("n", 1) completions = [] for _ in range(n): response = self.request(prompt, **kwargs) self.log_usage(response) # This is the original behavior in dspy/dsp/modules/anthropic.py. # Comment it out because it can cause "IndexError: list index out of range" silently # which is not transparent to developers. # if only_completed and response.stop_reason == "max_tokens": # continue completions = [c.text for c in response.content] return completions class VLLMClient(dspy.dsp.LM): """A client compatible with vLLM HTTP server. vLLM HTTP server is designed to be compatible with the OpenAI API. Use OpenAI client to interact with the server. """ def __init__( self, model, port, model_type: Literal["chat", "text"] = "text", url="http://localhost", api_key="null", **kwargs, ): """Check out https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html for more information.""" super().__init__(model=model) # Store additional kwargs for the generate method. self.kwargs = {**self.kwargs, **kwargs} self.model = model self.base_url = f"{url}:{port}/v1/" if model_type == "chat": self.base_url += "chat/" self.client = OpenAI(base_url=self.base_url, api_key=api_key) self.prompt_tokens = 0 self.completion_tokens = 0 self._token_usage_lock = threading.Lock() def basic_request(self, prompt, **kwargs): completion = self.client.chat.completions.create( **kwargs, messages=[{"role": "user", "content": prompt}], ) return completion @backoff.on_exception( backoff.expo, ERRORS, max_time=1000, on_backoff=backoff_hdlr, ) def request(self, prompt: str, **kwargs): return self.basic_request(prompt, **kwargs) def log_usage(self, response): """Log the total tokens from the response.""" usage_data = response.usage if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.prompt_tokens self.completion_tokens += usage_data.completion_tokens def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.kwargs.get("model") or self.kwargs.get("engine"): { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage def __call__(self, prompt: str, **kwargs): kwargs = {**self.kwargs, **kwargs} try: response = self.request(prompt, **kwargs) except Exception as e: print(f"Failed to generate completion: {e}") raise Exception(e) self.log_usage(response) choices = response.choices completions = [choice.message.content for choice in choices] history = { "prompt": prompt, "response": response, "kwargs": kwargs, } self.history.append(history) return completions class OllamaClient(dspy.OllamaLocal): """A wrapper class for dspy.OllamaClient.""" def __init__(self, model, port, url="http://localhost", **kwargs): """Copied from dspy/dsp/modules/hf_client.py with the addition of storing additional kwargs.""" # Check if the URL has 'http://' or 'https://' if not url.startswith("http://") and not url.startswith("https://"): url = "http://" + url super().__init__(model=model, base_url=f"{url}:{port}", **kwargs) # Store additional kwargs for the generate method. self.kwargs = {**self.kwargs, **kwargs} class TGIClient(dspy.HFClientTGI): def __init__(self, model, port, url, http_request_kwargs=None, **kwargs): super().__init__( model=model, port=port, url=url, http_request_kwargs=http_request_kwargs, **kwargs, ) def _generate(self, prompt, **kwargs): """Copied from dspy/dsp/modules/hf_client.py with the addition of removing hard-coded parameters.""" kwargs = {**self.kwargs, **kwargs} payload = { "inputs": prompt, "parameters": { "do_sample": kwargs["n"] > 1, "best_of": kwargs["n"], "details": kwargs["n"] > 1, **kwargs, }, } payload["parameters"] = openai_to_hf(**payload["parameters"]) # Comment out the following lines to remove the hard-coded parameters. # payload["parameters"]["temperature"] = max( # 0.1, payload["parameters"]["temperature"], # ) response = send_hftgi_request_v01_wrapped( f"{self.url}:{random.Random().choice(self.ports)}" + "/generate", url=self.url, ports=tuple(self.ports), json=payload, headers=self.headers, **self.http_request_kwargs, ) try: json_response = response.json() # completions = json_response["generated_text"] completions = [json_response["generated_text"]] if ( "details" in json_response and "best_of_sequences" in json_response["details"] ): completions += [ x["generated_text"] for x in json_response["details"]["best_of_sequences"] ] response = {"prompt": prompt, "choices": [{"text": c} for c in completions]} return response except Exception: print("Failed to parse JSON response:", response.text) raise Exception("Received invalid JSON response from server") class TogetherClient(dspy.HFModel): """A wrapper class for dspy.Together.""" def __init__( self, model, api_key: Optional[str] = None, apply_tokenizer_chat_template=False, hf_tokenizer_name=None, model_type: Literal["chat", "text"] = "chat", **kwargs, ): """Copied from dspy/dsp/modules/hf_client.py with the support of applying tokenizer chat template.""" super().__init__(model=model, is_client=True) self.session = requests.Session() self.api_key = api_key = ( os.environ.get("TOGETHER_API_KEY") if api_key is None else api_key ) self.model = model self.model_type = model_type if os.getenv("TOGETHER_API_BASE") is None: if self.model_type == "chat": self.api_base = "https://api.together.xyz/v1/chat/completions" else: self.api_base = "https://api.together.xyz/v1/completions" else: self.api_base = os.getenv("TOGETHER_API_BASE") # self.use_inst_template = False # if any(keyword in self.model.lower() for keyword in ["inst", "instruct"]): # self.use_inst_template = True self.apply_tokenizer_chat_template = apply_tokenizer_chat_template if self.apply_tokenizer_chat_template: logging.info("Loading huggingface tokenizer.") if hf_tokenizer_name is None: hf_tokenizer_name = self.model self.tokenizer = AutoTokenizer.from_pretrained( hf_tokenizer_name, cache_dir=kwargs.get("cache_dir", None) ) stop_default = "\n\n---" self.kwargs = { "temperature": kwargs.get("temperature", 0.0), "max_tokens": min(kwargs.get("max_tokens", 4096), 4096), "top_p": kwargs.get("top_p", 1.0), "top_k": kwargs.get("top_k", 1), "repetition_penalty": 1, "n": kwargs.pop("n", kwargs.pop("num_generations", 1)), "stop": stop_default if "stop" not in kwargs else kwargs["stop"], **kwargs, } self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 def log_usage(self, response): """Log the total tokens from the OpenAI API response.""" usage_data = response.get("usage") if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.get("prompt_tokens", 0) self.completion_tokens += usage_data.get("completion_tokens", 0) def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.model: { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage @backoff.on_exception( backoff.expo, ERRORS, max_time=1000, on_backoff=backoff_hdlr, ) def _generate(self, prompt, **kwargs): kwargs = {**self.kwargs, **kwargs} stop = kwargs.get("stop") temperature = kwargs.get("temperature") max_tokens = kwargs.get("max_tokens", 150) top_p = kwargs.get("top_p", 0.7) top_k = kwargs.get("top_k", 50) repetition_penalty = kwargs.get("repetition_penalty", 1) if self.apply_tokenizer_chat_template: prompt = self.tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], tokenize=False ) # prompt = f"[INST]{prompt}[/INST]" if self.use_inst_template else prompt if self.model_type == "chat": messages = [ { "role": "system", "content": "You are a helpful assistant. You must continue the user text directly without *any* additional interjections.", }, {"role": "user", "content": prompt}, ] body = { "model": self.model, "messages": messages, "temperature": temperature, "max_tokens": max_tokens, "top_p": top_p, "top_k": top_k, "repetition_penalty": repetition_penalty, "stop": stop, } else: body = { "model": self.model, "prompt": prompt, "temperature": temperature, "max_tokens": max_tokens, "top_p": top_p, "top_k": top_k, "repetition_penalty": repetition_penalty, "stop": stop, } headers = {"Authorization": f"Bearer {self.api_key}"} with self.session.post(self.api_base, headers=headers, json=body) as resp: resp_json = resp.json() # Log the token usage from the Together API response. self.log_usage(resp_json) if self.model_type == "chat": # completions = [resp_json['output'].get('choices', [])[0].get('message', {}).get('content', "")] completions = [ resp_json.get("choices", [])[0] .get("message", {}) .get("content", "") ] else: # completions = [resp_json['output'].get('choices', [])[0].get('text', "")] completions = [resp_json.get("choices", [])[0].get("text", "")] response = {"prompt": prompt, "choices": [{"text": c} for c in completions]} return response class GoogleModel(dspy.dsp.modules.lm.LM): """A wrapper class for Google Gemini API.""" def __init__( self, model: str, api_key: Optional[str] = None, **kwargs, ): """You can use `genai.list_models()` to get a list of available models.""" super().__init__(model) try: import google.generativeai as genai except ImportError as err: raise ImportError( "GoogleModel requires `pip install google-generativeai`." ) from err api_key = os.environ.get("GOOGLE_API_KEY") if api_key is None else api_key genai.configure(api_key=api_key) kwargs = { "candidate_count": 1, # Caveat: Gemini API supports only one candidate for now. "temperature": ( 0.0 if "temperature" not in kwargs else kwargs["temperature"] ), "max_output_tokens": kwargs["max_tokens"], "top_p": 1, "top_k": 1, **kwargs, } kwargs.pop("max_tokens", None) # GenerationConfig cannot accept max_tokens self.model = model self.config = genai.GenerationConfig(**kwargs) self.llm = genai.GenerativeModel( model_name=model, generation_config=self.config ) self.kwargs = { "n": 1, **kwargs, } self.history: list[dict[str, Any]] = [] self._token_usage_lock = threading.Lock() self.prompt_tokens = 0 self.completion_tokens = 0 def log_usage(self, response): """Log the total tokens from the Google API response.""" usage_data = response.usage_metadata if usage_data: with self._token_usage_lock: self.prompt_tokens += usage_data.prompt_token_count self.completion_tokens += usage_data.candidates_token_count def get_usage_and_reset(self): """Get the total tokens used and reset the token usage.""" usage = { self.model: { "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } } self.prompt_tokens = 0 self.completion_tokens = 0 return usage def basic_request(self, prompt: str, **kwargs): raw_kwargs = kwargs kwargs = { **self.kwargs, **kwargs, } # Google disallows "n" arguments. n = kwargs.pop("n", None) response = self.llm.generate_content(prompt, generation_config=kwargs) history = { "prompt": prompt, "response": [response.to_dict()], "kwargs": kwargs, "raw_kwargs": raw_kwargs, } self.history.append(history) return response @backoff.on_exception( backoff.expo, (Exception,), max_time=1000, max_tries=8, on_backoff=backoff_hdlr, giveup=giveup_hdlr, ) def request(self, prompt: str, **kwargs): """Handles retrieval of completions from Google whilst handling API errors""" return self.basic_request(prompt, **kwargs) def __call__( self, prompt: str, only_completed: bool = True, return_sorted: bool = False, **kwargs, ): assert only_completed, "for now" assert return_sorted is False, "for now" n = kwargs.pop("n", 1) completions = [] for _ in range(n): response = self.request(prompt, **kwargs) self.log_usage(response) completions.append(response.parts[0].text) return completions # ======================================================================== ================================================ FILE: knowledge_storm/logging_wrapper.py ================================================ from contextlib import contextmanager import time import pytz from datetime import datetime # Define California timezone CALIFORNIA_TZ = pytz.timezone("America/Los_Angeles") class EventLog: def __init__(self, event_name): self.event_name = event_name self.start_time = None self.end_time = None self.child_events = {} def record_start_time(self): self.start_time = datetime.now( pytz.utc ) # Store in UTC for consistent timezone conversion def record_end_time(self): self.end_time = datetime.now( pytz.utc ) # Store in UTC for consistent timezone conversion def get_total_time(self): if self.start_time and self.end_time: return (self.end_time - self.start_time).total_seconds() return 0 def get_start_time(self): if self.start_time: # Format to milliseconds return self.start_time.astimezone(CALIFORNIA_TZ).strftime( "%Y-%m-%d %H:%M:%S.%f" )[:-3] return None def get_end_time(self): if self.end_time: # Format to milliseconds return self.end_time.astimezone(CALIFORNIA_TZ).strftime( "%Y-%m-%d %H:%M:%S.%f" )[:-3] return None def add_child_event(self, child_event): self.child_events[child_event.event_name] = child_event def get_child_events(self): return self.child_events class LoggingWrapper: def __init__(self, lm_config): self.logging_dict = {} self.lm_config = lm_config self.current_pipeline_stage = None self.event_stack = [] self.pipeline_stage_active = False def _pipeline_stage_start(self, pipeline_stage: str): if self.pipeline_stage_active: raise RuntimeError( "A pipeline stage is already active. End the current stage before starting a new one." ) self.current_pipeline_stage = pipeline_stage self.logging_dict[pipeline_stage] = { "time_usage": {}, "lm_usage": {}, "lm_history": [], "query_count": 0, } self.pipeline_stage_active = True def _event_start(self, event_name: str): if not self.pipeline_stage_active: raise RuntimeError("No pipeline stage is currently active.") if not self.event_stack and self.current_pipeline_stage: # Top-level event (directly under the pipeline stage) if ( event_name not in self.logging_dict[self.current_pipeline_stage]["time_usage"] ): event = EventLog(event_name=event_name) event.record_start_time() self.logging_dict[self.current_pipeline_stage]["time_usage"][ event_name ] = event self.event_stack.append(event) else: self.logging_dict[self.current_pipeline_stage]["time_usage"][ event_name ].record_start_time() elif self.event_stack: # Nested event (under another event) parent_event = self.event_stack[-1] if event_name not in parent_event.get_child_events(): event = EventLog(event_name=event_name) event.record_start_time() parent_event.add_child_event(event) self.logging_dict[self.current_pipeline_stage]["time_usage"][ event_name ] = event self.event_stack.append(event) else: parent_event.get_child_events()[event_name].record_start_time() else: raise RuntimeError( "Cannot start an event without an active pipeline stage or parent event." ) def _event_end(self, event_name: str): if not self.pipeline_stage_active: raise RuntimeError("No pipeline stage is currently active.") if not self.event_stack: raise RuntimeError("No parent event is currently active.") if self.event_stack: current_event_log = self.event_stack[-1] if event_name in current_event_log.get_child_events(): current_event_log.get_child_events()[event_name].record_end_time() elif ( event_name in self.logging_dict[self.current_pipeline_stage]["time_usage"] ): self.logging_dict[self.current_pipeline_stage]["time_usage"][ event_name ].record_end_time() else: raise AssertionError( f"Failure to record end time for event {event_name}. Start time is not recorded." ) if current_event_log.event_name == event_name: self.event_stack.pop() else: raise RuntimeError("Cannot end an event without an active parent event.") def _pipeline_stage_end(self): if not self.pipeline_stage_active: raise RuntimeError("No pipeline stage is currently active to end.") self.logging_dict[self.current_pipeline_stage][ "lm_usage" ] = self.lm_config.collect_and_reset_lm_usage() self.logging_dict[self.current_pipeline_stage][ "lm_history" ] = self.lm_config.collect_and_reset_lm_history() self.pipeline_stage_active = False def add_query_count(self, count): if not self.pipeline_stage_active: raise RuntimeError( "No pipeline stage is currently active to add query count." ) self.logging_dict[self.current_pipeline_stage]["query_count"] += count @contextmanager def log_event(self, event_name): if not self.pipeline_stage_active: raise RuntimeError("No pipeline stage is currently active.") self._event_start(event_name) yield self._event_end(event_name) @contextmanager def log_pipeline_stage(self, pipeline_stage): if self.pipeline_stage_active: print( "A pipeline stage is already active, ending the current stage safely." ) self._pipeline_stage_end() start_time = time.time() try: self._pipeline_stage_start(pipeline_stage) yield except Exception as e: print(f"Error occurred during pipeline stage '{pipeline_stage}': {e}") finally: self.logging_dict[self.current_pipeline_stage]["total_wall_time"] = ( time.time() - start_time ) self._pipeline_stage_end() def dump_logging_and_reset(self, reset_logging=True): log_dump = {} for pipeline_stage, pipeline_log in self.logging_dict.items(): time_stamp_log = { event_name: { "total_time_seconds": event.get_total_time(), "start_time": event.get_start_time(), "end_time": event.get_end_time(), } for event_name, event in pipeline_log["time_usage"].items() } log_dump[pipeline_stage] = { "time_usage": time_stamp_log, "lm_usage": pipeline_log["lm_usage"], "lm_history": pipeline_log["lm_history"], "query_count": pipeline_log["query_count"], "total_wall_time": pipeline_log["total_wall_time"], } if reset_logging: self.logging_dict.clear() return log_dump ================================================ FILE: knowledge_storm/rm.py ================================================ import logging import os from typing import Callable, Union, List import backoff import dspy import requests from dsp import backoff_hdlr, giveup_hdlr from .utils import WebPageHelper class YouRM(dspy.Retrieve): def __init__(self, ydc_api_key=None, k=3, is_valid_source: Callable = None): super().__init__(k=k) if not ydc_api_key and not os.environ.get("YDC_API_KEY"): raise RuntimeError( "You must supply ydc_api_key or set environment variable YDC_API_KEY" ) elif ydc_api_key: self.ydc_api_key = ydc_api_key else: self.ydc_api_key = os.environ["YDC_API_KEY"] self.usage = 0 # If not None, is_valid_source shall be a function that takes a URL and returns a boolean. if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"YouRM": usage} def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search with You.com for self.k top passages for query or queries Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of urls to exclude from the search results. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) collected_results = [] for query in queries: try: headers = {"X-API-Key": self.ydc_api_key} results = requests.get( f"https://api.ydc-index.io/search?query={query}", headers=headers, ).json() authoritative_results = [] for r in results["hits"]: if self.is_valid_source(r["url"]) and r["url"] not in exclude_urls: authoritative_results.append(r) if "hits" in results: collected_results.extend(authoritative_results[: self.k]) except Exception as e: logging.error(f"Error occurs when searching query {query}: {e}") return collected_results class BingSearch(dspy.Retrieve): def __init__( self, bing_search_api_key=None, k=3, is_valid_source: Callable = None, min_char_count: int = 150, snippet_chunk_size: int = 1000, webpage_helper_max_threads=10, mkt="en-US", language="en", **kwargs, ): """ Params: min_char_count: Minimum character count for the article to be considered valid. snippet_chunk_size: Maximum character count for each snippet. webpage_helper_max_threads: Maximum number of threads to use for webpage helper. mkt, language, **kwargs: Bing search API parameters. - Reference: https://learn.microsoft.com/en-us/bing/search-apis/bing-web-search/reference/query-parameters """ super().__init__(k=k) if not bing_search_api_key and not os.environ.get("BING_SEARCH_API_KEY"): raise RuntimeError( "You must supply bing_search_subscription_key or set environment variable BING_SEARCH_API_KEY" ) elif bing_search_api_key: self.bing_api_key = bing_search_api_key else: self.bing_api_key = os.environ["BING_SEARCH_API_KEY"] self.endpoint = "https://api.bing.microsoft.com/v7.0/search" self.params = {"mkt": mkt, "setLang": language, "count": k, **kwargs} self.webpage_helper = WebPageHelper( min_char_count=min_char_count, snippet_chunk_size=snippet_chunk_size, max_thread_num=webpage_helper_max_threads, ) self.usage = 0 # If not None, is_valid_source shall be a function that takes a URL and returns a boolean. if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"BingSearch": usage} def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search with Bing for self.k top passages for query or queries Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of urls to exclude from the search results. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) url_to_results = {} headers = {"Ocp-Apim-Subscription-Key": self.bing_api_key} for query in queries: try: results = requests.get( self.endpoint, headers=headers, params={**self.params, "q": query} ).json() for d in results["webPages"]["value"]: if self.is_valid_source(d["url"]) and d["url"] not in exclude_urls: url_to_results[d["url"]] = { "url": d["url"], "title": d["name"], "description": d["snippet"], } except Exception as e: logging.error(f"Error occurs when searching query {query}: {e}") valid_url_to_snippets = self.webpage_helper.urls_to_snippets( list(url_to_results.keys()) ) collected_results = [] for url in valid_url_to_snippets: r = url_to_results[url] r["snippets"] = valid_url_to_snippets[url]["snippets"] collected_results.append(r) return collected_results class VectorRM(dspy.Retrieve): """Retrieve information from custom documents using Qdrant. To be compatible with STORM, the custom documents should have the following fields: - content: The main text content of the document. - title: The title of the document. - url: The URL of the document. STORM use url as the unique identifier of the document, so ensure different documents have different urls. - description (optional): The description of the document. The documents should be stored in a CSV file. """ def __init__( self, collection_name: str, embedding_model: str, device: str = "mps", k: int = 3, ): from langchain_huggingface import HuggingFaceEmbeddings """ Params: collection_name: Name of the Qdrant collection. embedding_model: Name of the Hugging Face embedding model. device: Device to run the embeddings model on, can be "mps", "cuda", "cpu". k: Number of top chunks to retrieve. """ super().__init__(k=k) self.usage = 0 # check if the collection is provided if not collection_name: raise ValueError("Please provide a collection name.") # check if the embedding model is provided if not embedding_model: raise ValueError("Please provide an embedding model.") model_kwargs = {"device": device} encode_kwargs = {"normalize_embeddings": True} self.model = HuggingFaceEmbeddings( model_name=embedding_model, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs, ) self.collection_name = collection_name self.client = None self.qdrant = None def _check_collection(self): from langchain_qdrant import Qdrant """ Check if the Qdrant collection exists and create it if it does not. """ if self.client is None: raise ValueError("Qdrant client is not initialized.") if self.client.collection_exists(collection_name=f"{self.collection_name}"): print( f"Collection {self.collection_name} exists. Loading the collection..." ) self.qdrant = Qdrant( client=self.client, collection_name=self.collection_name, embeddings=self.model, ) else: raise ValueError( f"Collection {self.collection_name} does not exist. Please create the collection first." ) def init_online_vector_db(self, url: str, api_key: str): from qdrant_client import QdrantClient """ Initialize the Qdrant client that is connected to an online vector store with the given URL and API key. Args: url (str): URL of the Qdrant server. api_key (str): API key for the Qdrant server. """ if api_key is None: if not os.getenv("QDRANT_API_KEY"): raise ValueError("Please provide an api key.") api_key = os.getenv("QDRANT_API_KEY") if url is None: raise ValueError("Please provide a url for the Qdrant server.") try: self.client = QdrantClient(url=url, api_key=api_key) self._check_collection() except Exception as e: raise ValueError(f"Error occurs when connecting to the server: {e}") def init_offline_vector_db(self, vector_store_path: str): from qdrant_client import QdrantClient """ Initialize the Qdrant client that is connected to an offline vector store with the given vector store folder path. Args: vector_store_path (str): Path to the vector store. """ if vector_store_path is None: raise ValueError("Please provide a folder path.") try: self.client = QdrantClient(path=vector_store_path) self._check_collection() except Exception as e: raise ValueError(f"Error occurs when loading the vector store: {e}") def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"VectorRM": usage} def get_vector_count(self): """ Get the count of vectors in the collection. Returns: int: Number of vectors in the collection. """ return self.qdrant.client.count(collection_name=self.collection_name) def forward(self, query_or_queries: Union[str, List[str]], exclude_urls: List[str]): """ Search in your data for self.k top passages for query or queries. Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): Dummy parameter to match the interface. Does not have any effect. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) collected_results = [] for query in queries: related_docs = self.qdrant.similarity_search_with_score(query, k=self.k) for i in range(len(related_docs)): doc = related_docs[i][0] collected_results.append( { "description": doc.metadata["description"], "snippets": [doc.page_content], "title": doc.metadata["title"], "url": doc.metadata["url"], } ) return collected_results class StanfordOvalArxivRM(dspy.Retrieve): """[Alpha] This retrieval class is for internal use only, not intended for the public.""" def __init__(self, endpoint, k=3, rerank=True): super().__init__(k=k) self.endpoint = endpoint self.usage = 0 self.rerank = rerank def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"StanfordOvalArxivRM": usage} def _retrieve(self, query: str): payload = {"query": query, "num_blocks": self.k, "rerank": self.rerank} response = requests.post( self.endpoint, json=payload, headers={"Content-Type": "application/json"} ) # Check if the request was successful if response.status_code == 200: response_data_list = response.json()[0]["results"] results = [] for response_data in response_data_list: result = { "title": response_data["document_title"], "url": response_data["url"], "snippets": [response_data["content"]], "description": response_data.get("description", "N/A"), "meta": { key: value for key, value in response_data.items() if key not in ["document_title", "url", "content"] }, } results.append(result) return results else: raise Exception( f"Error: Unable to retrieve results. Status code: {response.status_code}" ) def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): collected_results = [] queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) for query in queries: try: results = self._retrieve(query) collected_results.extend(results) except Exception as e: logging.error(f"Error occurs when searching query {query}: {e}") return collected_results class SerperRM(dspy.Retrieve): """Retrieve information from custom queries using Serper.dev.""" def __init__( self, serper_search_api_key=None, k=3, query_params=None, ENABLE_EXTRA_SNIPPET_EXTRACTION=False, min_char_count: int = 150, snippet_chunk_size: int = 1000, webpage_helper_max_threads=10, ): """Args: serper_search_api_key str: API key to run serper, can be found by creating an account on https://serper.dev/ query_params (dict or list of dict): parameters in dictionary or list of dictionaries that has a max size of 100 that will be used to query. Commonly used fields are as follows (see more information in https://serper.dev/playground): q str: query that will be used with google search type str: type that will be used for browsing google. Types are search, images, video, maps, places, etc. gl str: Country that will be focused on for the search location str: Country where the search will originate from. All locates can be found here: https://api.serper.dev/locations. autocorrect bool: Enable autocorrect on the queries while searching, if query is misspelled, will be updated. results int: Max number of results per page. page int: Max number of pages per call. tbs str: date time range, automatically set to any time by default. qdr:h str: Date time range for the past hour. qdr:d str: Date time range for the past 24 hours. qdr:w str: Date time range for past week. qdr:m str: Date time range for past month. qdr:y str: Date time range for past year. """ super().__init__(k=k) self.usage = 0 self.query_params = None self.ENABLE_EXTRA_SNIPPET_EXTRACTION = ENABLE_EXTRA_SNIPPET_EXTRACTION self.webpage_helper = WebPageHelper( min_char_count=min_char_count, snippet_chunk_size=snippet_chunk_size, max_thread_num=webpage_helper_max_threads, ) if query_params is None: self.query_params = {"num": k, "autocorrect": True, "page": 1} else: self.query_params = query_params self.query_params.update({"num": k}) self.serper_search_api_key = serper_search_api_key if not self.serper_search_api_key and not os.environ.get("SERPER_API_KEY"): raise RuntimeError( "You must supply a serper_search_api_key param or set environment variable SERPER_API_KEY" ) elif self.serper_search_api_key: self.serper_search_api_key = serper_search_api_key else: self.serper_search_api_key = os.environ["SERPER_API_KEY"] self.base_url = "https://google.serper.dev" def serper_runner(self, query_params): self.search_url = f"{self.base_url}/search" headers = { "X-API-KEY": self.serper_search_api_key, "Content-Type": "application/json", } response = requests.request( "POST", self.search_url, headers=headers, json=query_params ) if response == None: raise RuntimeError( f"Error had occurred while running the search process.\n Error is {response.reason}, had failed with status code {response.status_code}" ) return response.json() def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"SerperRM": usage} def forward(self, query_or_queries: Union[str, List[str]], exclude_urls: List[str]): """ Calls the API and searches for the query passed in. Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): Dummy parameter to match the interface. Does not have any effect. Returns: a list of dictionaries, each dictionary has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) self.results = [] collected_results = [] for query in queries: if query == "Queries:": continue query_params = self.query_params # All available parameters can be found in the playground: https://serper.dev/playground # Sets the json value for query to be the query that is being parsed. query_params["q"] = query # Sets the type to be search, can be images, video, places, maps etc that Google provides. query_params["type"] = "search" self.result = self.serper_runner(query_params) self.results.append(self.result) # Array of dictionaries that will be used by Storm to create the jsons collected_results = [] if self.ENABLE_EXTRA_SNIPPET_EXTRACTION: urls = [] for result in self.results: organic_results = result.get("organic", []) for organic in organic_results: url = organic.get("link") if url: urls.append(url) valid_url_to_snippets = self.webpage_helper.urls_to_snippets(urls) else: valid_url_to_snippets = {} for result in self.results: try: # An array of dictionaries that contains the snippets, title of the document and url that will be used. organic_results = result.get("organic") knowledge_graph = result.get("knowledgeGraph") for organic in organic_results: snippets = [organic.get("snippet")] if self.ENABLE_EXTRA_SNIPPET_EXTRACTION: snippets.extend( valid_url_to_snippets.get(url, {}).get("snippets", []) ) collected_results.append( { "snippets": snippets, "title": organic.get("title"), "url": organic.get("link"), "description": ( knowledge_graph.get("description") if knowledge_graph is not None else "" ), } ) except: continue return collected_results class BraveRM(dspy.Retrieve): def __init__( self, brave_search_api_key=None, k=3, is_valid_source: Callable = None ): super().__init__(k=k) if not brave_search_api_key and not os.environ.get("BRAVE_API_KEY"): raise RuntimeError( "You must supply brave_search_api_key or set environment variable BRAVE_API_KEY" ) elif brave_search_api_key: self.brave_search_api_key = brave_search_api_key else: self.brave_search_api_key = os.environ["BRAVE_API_KEY"] self.usage = 0 # If not None, is_valid_source shall be a function that takes a URL and returns a boolean. if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"BraveRM": usage} def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search with api.search.brave.com for self.k top passages for query or queries Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of urls to exclude from the search results. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) collected_results = [] for query in queries: try: headers = { "Accept": "application/json", "Accept-Encoding": "gzip", "X-Subscription-Token": self.brave_search_api_key, } response = requests.get( f"https://api.search.brave.com/res/v1/web/search?result_filter=web&q={query}", headers=headers, ).json() results = response.get("web", {}).get("results", []) for result in results: collected_results.append( { "snippets": result.get("extra_snippets", []), "title": result.get("title"), "url": result.get("url"), "description": result.get("description"), } ) except Exception as e: logging.error(f"Error occurs when searching query {query}: {e}") return collected_results class SearXNG(dspy.Retrieve): def __init__( self, searxng_api_url, searxng_api_key=None, k=3, is_valid_source: Callable = None, ): """Initialize the SearXNG search retriever. Please set up SearXNG according to https://docs.searxng.org/index.html. Args: searxng_api_url (str): The URL of the SearXNG API. Consult SearXNG documentation for details. searxng_api_key (str, optional): The API key for the SearXNG API. Defaults to None. Consult SearXNG documentation for details. k (int, optional): The number of top passages to retrieve. Defaults to 3. is_valid_source (Callable, optional): A function that takes a URL and returns a boolean indicating if the source is valid. Defaults to None. """ super().__init__(k=k) if not searxng_api_url: raise RuntimeError("You must supply searxng_api_url") self.searxng_api_url = searxng_api_url self.searxng_api_key = searxng_api_key self.usage = 0 if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"SearXNG": usage} def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search with SearxNG for self.k top passages for query or queries Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of urls to exclude from the search results. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) collected_results = [] headers = ( {"Authorization": f"Bearer {self.searxng_api_key}"} if self.searxng_api_key else {} ) for query in queries: try: params = {"q": query, "format": "json"} response = requests.get( self.searxng_api_url, headers=headers, params=params ) results = response.json() for r in results["results"]: if self.is_valid_source(r["url"]) and r["url"] not in exclude_urls: collected_results.append( { "description": r.get("content", ""), "snippets": [r.get("content", "")], "title": r.get("title", ""), "url": r["url"], } ) except Exception as e: logging.error(f"Error occurs when searching query {query}: {e}") return collected_results class DuckDuckGoSearchRM(dspy.Retrieve): """Retrieve information from custom queries using DuckDuckGo.""" def __init__( self, k: int = 3, is_valid_source: Callable = None, min_char_count: int = 150, snippet_chunk_size: int = 1000, webpage_helper_max_threads=10, safe_search: str = "On", region: str = "us-en", ): """ Params: min_char_count: Minimum character count for the article to be considered valid. snippet_chunk_size: Maximum character count for each snippet. webpage_helper_max_threads: Maximum number of threads to use for webpage helper. **kwargs: Additional parameters for the OpenAI API. """ super().__init__(k=k) try: from duckduckgo_search import DDGS except ImportError as err: raise ImportError( "Duckduckgo requires `pip install duckduckgo_search`." ) from err self.k = k self.webpage_helper = WebPageHelper( min_char_count=min_char_count, snippet_chunk_size=snippet_chunk_size, max_thread_num=webpage_helper_max_threads, ) self.usage = 0 # All params for search can be found here: # https://duckduckgo.com/duckduckgo-help-pages/settings/params/ # Sets the backend to be api self.duck_duck_go_backend = "api" # Only gets safe search results self.duck_duck_go_safe_search = safe_search # Specifies the region that the search will use self.duck_duck_go_region = region # If not None, is_valid_source shall be a function that takes a URL and returns a boolean. if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True # Import the duckduckgo search library found here: https://github.com/deedy5/duckduckgo_search self.ddgs = DDGS() def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"DuckDuckGoRM": usage} @backoff.on_exception( backoff.expo, (Exception,), max_time=1000, max_tries=8, on_backoff=backoff_hdlr, giveup=giveup_hdlr, ) def request(self, query: str): results = self.ddgs.text( query, max_results=self.k, backend=self.duck_duck_go_backend ) return results def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search with DuckDuckGoSearch for self.k top passages for query or queries Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of urls to exclude from the search results. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) collected_results = [] for query in queries: # list of dicts that will be parsed to return results = self.request(query) for d in results: # assert d is dict if not isinstance(d, dict): print(f"Invalid result: {d}\n") continue try: # ensure keys are present url = d.get("href", None) title = d.get("title", None) description = d.get("description", title) snippets = [d.get("body", None)] # raise exception of missing key(s) if not all([url, title, description, snippets]): raise ValueError(f"Missing key(s) in result: {d}") if self.is_valid_source(url) and url not in exclude_urls: result = { "url": url, "title": title, "description": description, "snippets": snippets, } collected_results.append(result) else: print(f"invalid source {url} or url in exclude_urls") except Exception as e: print(f"Error occurs when processing {result=}: {e}\n") print(f"Error occurs when searching query {query}: {e}") return collected_results class TavilySearchRM(dspy.Retrieve): """Retrieve information from custom queries using Tavily. Documentation and examples can be found at https://docs.tavily.com/docs/python-sdk/tavily-search/examples""" def __init__( self, tavily_search_api_key=None, k: int = 3, is_valid_source: Callable = None, min_char_count: int = 150, snippet_chunk_size: int = 1000, webpage_helper_max_threads=10, include_raw_content=False, ): """ Params: tavily_search_api_key str: API key for tavily that can be retrieved from https://tavily.com/ min_char_count: Minimum character count for the article to be considered valid. snippet_chunk_size: Maximum character count for each snippet. webpage_helper_max_threads: Maximum number of threads to use for webpage helper. include_raw_content bool: Boolean that is used to determine if the full text should be returned. """ super().__init__(k=k) try: from tavily import TavilyClient except ImportError as err: raise ImportError("Tavily requires `pip install tavily-python`.") from err if not tavily_search_api_key and not os.environ.get("TAVILY_API_KEY"): raise RuntimeError( "You must supply tavily_search_api_key or set environment variable TAVILY_API_KEY" ) elif tavily_search_api_key: self.tavily_search_api_key = tavily_search_api_key else: self.tavily_search_api_key = os.environ["TAVILY_API_KEY"] self.k = k self.webpage_helper = WebPageHelper( min_char_count=min_char_count, snippet_chunk_size=snippet_chunk_size, max_thread_num=webpage_helper_max_threads, ) self.usage = 0 # Creates client instance that will use search. Full search params are here: # https://docs.tavily.com/docs/python-sdk/tavily-search/examples self.tavily_client = TavilyClient(api_key=self.tavily_search_api_key) self.include_raw_content = include_raw_content # If not None, is_valid_source shall be a function that takes a URL and returns a boolean. if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"TavilySearchRM": usage} def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search with TavilySearch for self.k top passages for query or queries Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of urls to exclude from the search results. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) collected_results = [] for query in queries: args = { "max_results": self.k, "include_raw_contents": self.include_raw_content, } # list of dicts that will be parsed to return responseData = self.tavily_client.search(query) results = responseData.get("results") for d in results: # assert d is dict if not isinstance(d, dict): print(f"Invalid result: {d}\n") continue try: # ensure keys are present url = d.get("url", None) title = d.get("title", None) description = d.get("content", None) snippets = [] if d.get("raw_body_content"): snippets.append(d.get("raw_body_content")) else: snippets.append(d.get("content")) # raise exception of missing key(s) if not all([url, title, description, snippets]): raise ValueError(f"Missing key(s) in result: {d}") if self.is_valid_source(url) and url not in exclude_urls: result = { "url": url, "title": title, "description": description, "snippets": snippets, } collected_results.append(result) else: print(f"invalid source {url} or url in exclude_urls") except Exception as e: print(f"Error occurs when processing {result=}: {e}\n") print(f"Error occurs when searching query {query}: {e}") return collected_results class GoogleSearch(dspy.Retrieve): def __init__( self, google_search_api_key=None, google_cse_id=None, k=3, is_valid_source: Callable = None, min_char_count: int = 150, snippet_chunk_size: int = 1000, webpage_helper_max_threads=10, ): """ Params: google_search_api_key: Google API key. Check out https://developers.google.com/custom-search/v1/overview "API key" section google_cse_id: Custom search engine ID. Check out https://developers.google.com/custom-search/v1/overview "Search engine ID" section k: Number of top results to retrieve. is_valid_source: Optional function to filter valid sources. min_char_count: Minimum character count for the article to be considered valid. snippet_chunk_size: Maximum character count for each snippet. webpage_helper_max_threads: Maximum number of threads to use for webpage helper. """ super().__init__(k=k) try: from googleapiclient.discovery import build except ImportError as err: raise ImportError( "GoogleSearch requires `pip install google-api-python-client`." ) from err if not google_search_api_key and not os.environ.get("GOOGLE_SEARCH_API_KEY"): raise RuntimeError( "You must supply google_search_api_key or set the GOOGLE_SEARCH_API_KEY environment variable" ) if not google_cse_id and not os.environ.get("GOOGLE_CSE_ID"): raise RuntimeError( "You must supply google_cse_id or set the GOOGLE_CSE_ID environment variable" ) self.google_search_api_key = ( google_search_api_key or os.environ["GOOGLE_SEARCH_API_KEY"] ) self.google_cse_id = google_cse_id or os.environ["GOOGLE_CSE_ID"] if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True self.service = build( "customsearch", "v1", developerKey=self.google_search_api_key ) self.webpage_helper = WebPageHelper( min_char_count=min_char_count, snippet_chunk_size=snippet_chunk_size, max_thread_num=webpage_helper_max_threads, ) self.usage = 0 def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"GoogleSearch": usage} def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search using Google Custom Search API for self.k top results for query or queries. Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of URLs to exclude from the search results. Returns: A list of dicts, each dict has keys: 'title', 'url', 'snippet', 'description'. """ queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) url_to_results = {} for query in queries: try: response = ( self.service.cse() .list( q=query, cx=self.google_cse_id, num=self.k, ) .execute() ) for item in response.get("items", []): if ( self.is_valid_source(item["link"]) and item["link"] not in exclude_urls ): url_to_results[item["link"]] = { "title": item["title"], "url": item["link"], # "snippet": item.get("snippet", ""), # Google search snippet is very short. "description": item.get("snippet", ""), } except Exception as e: logging.error(f"Error occurred while searching query {query}: {e}") valid_url_to_snippets = self.webpage_helper.urls_to_snippets( list(url_to_results.keys()) ) collected_results = [] for url in valid_url_to_snippets: r = url_to_results[url] r["snippets"] = valid_url_to_snippets[url]["snippets"] collected_results.append(r) return collected_results class AzureAISearch(dspy.Retrieve): """Retrieve information from custom queries using Azure AI Search. General Documentation: https://learn.microsoft.com/en-us/azure/search/search-create-service-portal. Python Documentation: https://learn.microsoft.com/en-us/python/api/overview/azure/search-documents-readme?view=azure-python. """ def __init__( self, azure_ai_search_api_key=None, azure_ai_search_url=None, azure_ai_search_index_name=None, k=3, is_valid_source: Callable = None, ): """ Params: azure_ai_search_api_key: Azure AI Search API key. Check out https://learn.microsoft.com/en-us/azure/search/search-security-api-keys?tabs=rest-use%2Cportal-find%2Cportal-query "API key" section azure_ai_search_url: Custom Azure AI Search Endpoint URL. Check out https://learn.microsoft.com/en-us/azure/search/search-create-service-portal#name-the-service azure_ai_search_index_name: Custom Azure AI Search Index Name. Check out https://learn.microsoft.com/en-us/azure/search/search-how-to-create-search-index?tabs=portal k: Number of top results to retrieve. is_valid_source: Optional function to filter valid sources. min_char_count: Minimum character count for the article to be considered valid. snippet_chunk_size: Maximum character count for each snippet. webpage_helper_max_threads: Maximum number of threads to use for webpage helper. """ super().__init__(k=k) try: from azure.core.credentials import AzureKeyCredential from azure.search.documents import SearchClient except ImportError as err: raise ImportError( "AzureAISearch requires `pip install azure-search-documents`." ) from err if not azure_ai_search_api_key and not os.environ.get( "AZURE_AI_SEARCH_API_KEY" ): raise RuntimeError( "You must supply azure_ai_search_api_key or set environment variable AZURE_AI_SEARCH_API_KEY" ) elif azure_ai_search_api_key: self.azure_ai_search_api_key = azure_ai_search_api_key else: self.azure_ai_search_api_key = os.environ["AZURE_AI_SEARCH_API_KEY"] if not azure_ai_search_url and not os.environ.get("AZURE_AI_SEARCH_URL"): raise RuntimeError( "You must supply azure_ai_search_url or set environment variable AZURE_AI_SEARCH_URL" ) elif azure_ai_search_url: self.azure_ai_search_url = azure_ai_search_url else: self.azure_ai_search_url = os.environ["AZURE_AI_SEARCH_URL"] if not azure_ai_search_index_name and not os.environ.get( "AZURE_AI_SEARCH_INDEX_NAME" ): raise RuntimeError( "You must supply azure_ai_search_index_name or set environment variable AZURE_AI_SEARCH_INDEX_NAME" ) elif azure_ai_search_index_name: self.azure_ai_search_index_name = azure_ai_search_index_name else: self.azure_ai_search_index_name = os.environ["AZURE_AI_SEARCH_INDEX_NAME"] self.usage = 0 # If not None, is_valid_source shall be a function that takes a URL and returns a boolean. if is_valid_source: self.is_valid_source = is_valid_source else: self.is_valid_source = lambda x: True def get_usage_and_reset(self): usage = self.usage self.usage = 0 return {"AzureAISearch": usage} def forward( self, query_or_queries: Union[str, List[str]], exclude_urls: List[str] = [] ): """Search with Azure Open AI for self.k top passages for query or queries Args: query_or_queries (Union[str, List[str]]): The query or queries to search for. exclude_urls (List[str]): A list of urls to exclude from the search results. Returns: a list of Dicts, each dict has keys of 'description', 'snippets' (list of strings), 'title', 'url' """ try: from azure.core.credentials import AzureKeyCredential from azure.search.documents import SearchClient except ImportError as err: raise ImportError( "AzureAISearch requires `pip install azure-search-documents`." ) from err queries = ( [query_or_queries] if isinstance(query_or_queries, str) else query_or_queries ) self.usage += len(queries) collected_results = [] client = SearchClient( self.azure_ai_search_url, self.azure_ai_search_index_name, AzureKeyCredential(self.azure_ai_search_api_key), ) for query in queries: try: # https://learn.microsoft.com/en-us/python/api/azure-search-documents/azure.search.documents.searchclient?view=azure-python#azure-search-documents-searchclient-search results = client.search(search_text=query, top=1) for result in results: document = { "url": result["metadata_storage_path"], "title": result["title"], "description": "N/A", "snippets": [result["chunk"]], } collected_results.append(document) except Exception as e: logging.error(f"Error occurs when searching query {query}: {e}") return collected_results ================================================ FILE: knowledge_storm/storm_wiki/__init__.py ================================================ from .engine import * from .modules import * ================================================ FILE: knowledge_storm/storm_wiki/engine.py ================================================ import json import logging import os from dataclasses import dataclass, field from typing import Union, Literal, Optional import dspy from .modules.article_generation import StormArticleGenerationModule from .modules.article_polish import StormArticlePolishingModule from .modules.callback import BaseCallbackHandler from .modules.knowledge_curation import StormKnowledgeCurationModule from .modules.outline_generation import StormOutlineGenerationModule from .modules.persona_generator import StormPersonaGenerator from .modules.storm_dataclass import StormInformationTable, StormArticle from ..interface import Engine, LMConfigs, Retriever from ..lm import LitellmModel from ..utils import FileIOHelper, makeStringRed, truncate_filename class STORMWikiLMConfigs(LMConfigs): """Configurations for LLM used in different parts of STORM. Given that different parts in STORM framework have different complexity, we use different LLM configurations to achieve a balance between quality and efficiency. If no specific configuration is provided, we use the default setup in the paper. """ def __init__(self): self.conv_simulator_lm = ( None # LLM used in conversation simulator except for question asking. ) self.question_asker_lm = None # LLM used in question asking. self.outline_gen_lm = None # LLM used in outline generation. self.article_gen_lm = None # LLM used in article generation. self.article_polish_lm = None # LLM used in article polishing. def init_openai_model( self, openai_api_key: str, azure_api_key: str, openai_type: Literal["openai", "azure"], api_base: Optional[str] = None, api_version: Optional[str] = None, temperature: Optional[float] = 1.0, top_p: Optional[float] = 0.9, ): """Legacy: Corresponding to the original setup in the NAACL'24 paper.""" azure_kwargs = { "api_key": azure_api_key, "temperature": temperature, "top_p": top_p, "api_base": api_base, "api_version": api_version, } openai_kwargs = { "api_key": openai_api_key, "temperature": temperature, "top_p": top_p, "api_base": None, } if openai_type and openai_type == "openai": self.conv_simulator_lm = LitellmModel( model="gpt-4o-mini-2024-07-18", max_tokens=500, **openai_kwargs ) self.question_asker_lm = LitellmModel( model="gpt-4o-mini-2024-07-18", max_tokens=500, **openai_kwargs ) # 1/12/2024: Update gpt-4 to gpt-4-1106-preview. (Currently keep the original setup when using azure.) self.outline_gen_lm = LitellmModel( model="gpt-4-0125-preview", max_tokens=400, **openai_kwargs ) self.article_gen_lm = LitellmModel( model="gpt-4o-2024-05-13", max_tokens=700, **openai_kwargs ) self.article_polish_lm = LitellmModel( model="gpt-4o-2024-05-13", max_tokens=4000, **openai_kwargs ) elif openai_type and openai_type == "azure": self.conv_simulator_lm = LitellmModel( model="azure/gpt-4o-mini-2024-07-18", max_tokens=500, **openai_kwargs ) self.question_asker_lm = LitellmModel( model="azure/gpt-4o-mini-2024-07-18", max_tokens=500, **azure_kwargs, model_type="chat", ) # use combination of openai and azure-openai as azure-openai does not support gpt-4 in standard deployment self.outline_gen_lm = LitellmModel( model="azure/gpt-4o", max_tokens=400, **azure_kwargs, model_type="chat" ) self.article_gen_lm = LitellmModel( model="azure/gpt-4o-mini-2024-07-18", max_tokens=700, **azure_kwargs, model_type="chat", ) self.article_polish_lm = LitellmModel( model="azure/gpt-4o-mini-2024-07-18", max_tokens=4000, **azure_kwargs, model_type="chat", ) else: logging.warning( "No valid OpenAI API provider is provided. Cannot use default LLM configurations." ) def set_conv_simulator_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.conv_simulator_lm = model def set_question_asker_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.question_asker_lm = model def set_outline_gen_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.outline_gen_lm = model def set_article_gen_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.article_gen_lm = model def set_article_polish_lm(self, model: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.article_polish_lm = model @dataclass class STORMWikiRunnerArguments: """Arguments for controlling the STORM Wiki pipeline.""" output_dir: str = field( metadata={"help": "Output directory for the results."}, ) max_conv_turn: int = field( default=3, metadata={ "help": "Maximum number of questions in conversational question asking." }, ) max_perspective: int = field( default=3, metadata={ "help": "Maximum number of perspectives to consider in perspective-guided question asking." }, ) max_search_queries_per_turn: int = field( default=3, metadata={"help": "Maximum number of search queries to consider in each turn."}, ) disable_perspective: bool = field( default=False, metadata={"help": "If True, disable perspective-guided question asking."}, ) search_top_k: int = field( default=3, metadata={"help": "Top k search results to consider for each search query."}, ) retrieve_top_k: int = field( default=3, metadata={"help": "Top k collected references for each section title."}, ) max_thread_num: int = field( default=10, metadata={ "help": "Maximum number of threads to use. " "Consider reducing it if keep getting 'Exceed rate limit' error when calling LM API." }, ) class STORMWikiRunner(Engine): """STORM Wiki pipeline runner.""" def __init__( self, args: STORMWikiRunnerArguments, lm_configs: STORMWikiLMConfigs, rm ): super().__init__(lm_configs=lm_configs) self.args = args self.lm_configs = lm_configs self.retriever = Retriever(rm=rm, max_thread=self.args.max_thread_num) storm_persona_generator = StormPersonaGenerator( self.lm_configs.question_asker_lm ) self.storm_knowledge_curation_module = StormKnowledgeCurationModule( retriever=self.retriever, persona_generator=storm_persona_generator, conv_simulator_lm=self.lm_configs.conv_simulator_lm, question_asker_lm=self.lm_configs.question_asker_lm, max_search_queries_per_turn=self.args.max_search_queries_per_turn, search_top_k=self.args.search_top_k, max_conv_turn=self.args.max_conv_turn, max_thread_num=self.args.max_thread_num, ) self.storm_outline_generation_module = StormOutlineGenerationModule( outline_gen_lm=self.lm_configs.outline_gen_lm ) self.storm_article_generation = StormArticleGenerationModule( article_gen_lm=self.lm_configs.article_gen_lm, retrieve_top_k=self.args.retrieve_top_k, max_thread_num=self.args.max_thread_num, ) self.storm_article_polishing_module = StormArticlePolishingModule( article_gen_lm=self.lm_configs.article_gen_lm, article_polish_lm=self.lm_configs.article_polish_lm, ) self.lm_configs.init_check() self.apply_decorators() def run_knowledge_curation_module( self, ground_truth_url: str = "None", callback_handler: BaseCallbackHandler = None, ) -> StormInformationTable: ( information_table, conversation_log, ) = self.storm_knowledge_curation_module.research( topic=self.topic, ground_truth_url=ground_truth_url, callback_handler=callback_handler, max_perspective=self.args.max_perspective, disable_perspective=False, return_conversation_log=True, ) FileIOHelper.dump_json( conversation_log, os.path.join(self.article_output_dir, "conversation_log.json"), ) information_table.dump_url_to_info( os.path.join(self.article_output_dir, "raw_search_results.json") ) return information_table def run_outline_generation_module( self, information_table: StormInformationTable, callback_handler: BaseCallbackHandler = None, ) -> StormArticle: outline, draft_outline = self.storm_outline_generation_module.generate_outline( topic=self.topic, information_table=information_table, return_draft_outline=True, callback_handler=callback_handler, ) outline.dump_outline_to_file( os.path.join(self.article_output_dir, "storm_gen_outline.txt") ) draft_outline.dump_outline_to_file( os.path.join(self.article_output_dir, "direct_gen_outline.txt") ) return outline def run_article_generation_module( self, outline: StormArticle, information_table=StormInformationTable, callback_handler: BaseCallbackHandler = None, ) -> StormArticle: draft_article = self.storm_article_generation.generate_article( topic=self.topic, information_table=information_table, article_with_outline=outline, callback_handler=callback_handler, ) draft_article.dump_article_as_plain_text( os.path.join(self.article_output_dir, "storm_gen_article.txt") ) draft_article.dump_reference_to_file( os.path.join(self.article_output_dir, "url_to_info.json") ) return draft_article def run_article_polishing_module( self, draft_article: StormArticle, remove_duplicate: bool = False ) -> StormArticle: polished_article = self.storm_article_polishing_module.polish_article( topic=self.topic, draft_article=draft_article, remove_duplicate=remove_duplicate, ) FileIOHelper.write_str( polished_article.to_string(), os.path.join(self.article_output_dir, "storm_gen_article_polished.txt"), ) return polished_article def post_run(self): """ Post-run operations, including: 1. Dumping the run configuration. 2. Dumping the LLM call history. """ config_log = self.lm_configs.log() FileIOHelper.dump_json( config_log, os.path.join(self.article_output_dir, "run_config.json") ) llm_call_history = self.lm_configs.collect_and_reset_lm_history() with open( os.path.join(self.article_output_dir, "llm_call_history.jsonl"), "w" ) as f: for call in llm_call_history: if "kwargs" in call: call.pop( "kwargs" ) # All kwargs are dumped together to run_config.json. f.write(json.dumps(call) + "\n") def _load_information_table_from_local_fs(self, information_table_local_path): assert os.path.exists(information_table_local_path), makeStringRed( f"{information_table_local_path} not exists. Please set --do-research argument to prepare the conversation_log.json for this topic." ) return StormInformationTable.from_conversation_log_file( information_table_local_path ) def _load_outline_from_local_fs(self, topic, outline_local_path): assert os.path.exists(outline_local_path), makeStringRed( f"{outline_local_path} not exists. Please set --do-generate-outline argument to prepare the storm_gen_outline.txt for this topic." ) return StormArticle.from_outline_file(topic=topic, file_path=outline_local_path) def _load_draft_article_from_local_fs( self, topic, draft_article_path, url_to_info_path ): assert os.path.exists(draft_article_path), makeStringRed( f"{draft_article_path} not exists. Please set --do-generate-article argument to prepare the storm_gen_article.txt for this topic." ) assert os.path.exists(url_to_info_path), makeStringRed( f"{url_to_info_path} not exists. Please set --do-generate-article argument to prepare the url_to_info.json for this topic." ) article_text = FileIOHelper.load_str(draft_article_path) references = FileIOHelper.load_json(url_to_info_path) return StormArticle.from_string( topic_name=topic, article_text=article_text, references=references ) def run( self, topic: str, ground_truth_url: str = "", do_research: bool = True, do_generate_outline: bool = True, do_generate_article: bool = True, do_polish_article: bool = True, remove_duplicate: bool = False, callback_handler: BaseCallbackHandler = BaseCallbackHandler(), ): """ Run the STORM pipeline. Args: topic: The topic to research. ground_truth_url: A ground truth URL including a curated article about the topic. The URL will be excluded. do_research: If True, research the topic through information-seeking conversation; if False, expect conversation_log.json and raw_search_results.json to exist in the output directory. do_generate_outline: If True, generate an outline for the topic; if False, expect storm_gen_outline.txt to exist in the output directory. do_generate_article: If True, generate a curated article for the topic; if False, expect storm_gen_article.txt to exist in the output directory. do_polish_article: If True, polish the article by adding a summarization section and (optionally) removing duplicated content. remove_duplicate: If True, remove duplicated content. callback_handler: A callback handler to handle the intermediate results. """ assert ( do_research or do_generate_outline or do_generate_article or do_polish_article ), makeStringRed( "No action is specified. Please set at least one of --do-research, --do-generate-outline, --do-generate-article, --do-polish-article" ) self.topic = topic self.article_dir_name = truncate_filename( topic.replace(" ", "_").replace("/", "_") ) self.article_output_dir = os.path.join( self.args.output_dir, self.article_dir_name ) os.makedirs(self.article_output_dir, exist_ok=True) # research module information_table: StormInformationTable = None if do_research: information_table = self.run_knowledge_curation_module( ground_truth_url=ground_truth_url, callback_handler=callback_handler ) # outline generation module outline: StormArticle = None if do_generate_outline: # load information table if it's not initialized if information_table is None: information_table = self._load_information_table_from_local_fs( os.path.join(self.article_output_dir, "conversation_log.json") ) outline = self.run_outline_generation_module( information_table=information_table, callback_handler=callback_handler ) # article generation module draft_article: StormArticle = None if do_generate_article: if information_table is None: information_table = self._load_information_table_from_local_fs( os.path.join(self.article_output_dir, "conversation_log.json") ) if outline is None: outline = self._load_outline_from_local_fs( topic=topic, outline_local_path=os.path.join( self.article_output_dir, "storm_gen_outline.txt" ), ) draft_article = self.run_article_generation_module( outline=outline, information_table=information_table, callback_handler=callback_handler, ) # article polishing module if do_polish_article: if draft_article is None: draft_article_path = os.path.join( self.article_output_dir, "storm_gen_article.txt" ) url_to_info_path = os.path.join( self.article_output_dir, "url_to_info.json" ) draft_article = self._load_draft_article_from_local_fs( topic=topic, draft_article_path=draft_article_path, url_to_info_path=url_to_info_path, ) self.run_article_polishing_module( draft_article=draft_article, remove_duplicate=remove_duplicate ) ================================================ FILE: knowledge_storm/storm_wiki/modules/__init__.py ================================================ from .knowledge_curation import * from .persona_generator import * from .retriever import * from .storm_dataclass import * ================================================ FILE: knowledge_storm/storm_wiki/modules/article_generation.py ================================================ import concurrent.futures import copy import logging from concurrent.futures import as_completed from typing import List, Union import dspy from .callback import BaseCallbackHandler from .storm_dataclass import StormInformationTable, StormArticle from ...interface import ArticleGenerationModule, Information from ...utils import ArticleTextProcessing class StormArticleGenerationModule(ArticleGenerationModule): """ The interface for article generation stage. Given topic, collected information from knowledge curation stage, generated outline from outline generation stage, """ def __init__( self, article_gen_lm=Union[dspy.dsp.LM, dspy.dsp.HFModel], retrieve_top_k: int = 5, max_thread_num: int = 10, ): super().__init__() self.retrieve_top_k = retrieve_top_k self.article_gen_lm = article_gen_lm self.max_thread_num = max_thread_num self.section_gen = ConvToSection(engine=self.article_gen_lm) def generate_section( self, topic, section_name, information_table, section_outline, section_query ): collected_info: List[Information] = [] if information_table is not None: collected_info = information_table.retrieve_information( queries=section_query, search_top_k=self.retrieve_top_k ) output = self.section_gen( topic=topic, outline=section_outline, section=section_name, collected_info=collected_info, ) return { "section_name": section_name, "section_content": output.section, "collected_info": collected_info, } def generate_article( self, topic: str, information_table: StormInformationTable, article_with_outline: StormArticle, callback_handler: BaseCallbackHandler = None, ) -> StormArticle: """ Generate article for the topic based on the information table and article outline. Args: topic (str): The topic of the article. information_table (StormInformationTable): The information table containing the collected information. article_with_outline (StormArticle): The article with specified outline. callback_handler (BaseCallbackHandler): An optional callback handler that can be used to trigger custom callbacks at various stages of the article generation process. Defaults to None. """ information_table.prepare_table_for_retrieval() if article_with_outline is None: article_with_outline = StormArticle(topic_name=topic) sections_to_write = article_with_outline.get_first_level_section_names() section_output_dict_collection = [] if len(sections_to_write) == 0: logging.error( f"No outline for {topic}. Will directly search with the topic." ) section_output_dict = self.generate_section( topic=topic, section_name=topic, information_table=information_table, section_outline="", section_query=[topic], ) section_output_dict_collection = [section_output_dict] else: with concurrent.futures.ThreadPoolExecutor( max_workers=self.max_thread_num ) as executor: future_to_sec_title = {} for section_title in sections_to_write: # We don't want to write a separate introduction section. if section_title.lower().strip() == "introduction": continue # We don't want to write a separate conclusion section. if section_title.lower().strip().startswith( "conclusion" ) or section_title.lower().strip().startswith("summary"): continue section_query = article_with_outline.get_outline_as_list( root_section_name=section_title, add_hashtags=False ) queries_with_hashtags = article_with_outline.get_outline_as_list( root_section_name=section_title, add_hashtags=True ) section_outline = "\n".join(queries_with_hashtags) future_to_sec_title[ executor.submit( self.generate_section, topic, section_title, information_table, section_outline, section_query, ) ] = section_title for future in as_completed(future_to_sec_title): section_output_dict_collection.append(future.result()) article = copy.deepcopy(article_with_outline) for section_output_dict in section_output_dict_collection: article.update_section( parent_section_name=topic, current_section_content=section_output_dict["section_content"], current_section_info_list=section_output_dict["collected_info"], ) article.post_processing() return article class ConvToSection(dspy.Module): """Use the information collected from the information-seeking conversation to write a section.""" def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): super().__init__() self.write_section = dspy.Predict(WriteSection) self.engine = engine def forward( self, topic: str, outline: str, section: str, collected_info: List[Information] ): info = "" for idx, storm_info in enumerate(collected_info): info += f"[{idx + 1}]\n" + "\n".join(storm_info.snippets) info += "\n\n" info = ArticleTextProcessing.limit_word_count_preserve_newline(info, 1500) with dspy.settings.context(lm=self.engine): section = ArticleTextProcessing.clean_up_section( self.write_section(topic=topic, info=info, section=section).output ) return dspy.Prediction(section=section) class WriteSection(dspy.Signature): """Write a Wikipedia section based on the collected information. Here is the format of your writing: 1. Use "#" Title" to indicate section title, "##" Title" to indicate subsection title, "###" Title" to indicate subsubsection title, and so on. 2. Use [1], [2], ..., [n] in line (for example, "The capital of the United States is Washington, D.C.[1][3]."). You DO NOT need to include a References or Sources section to list the sources at the end. """ info = dspy.InputField(prefix="The collected information:\n", format=str) topic = dspy.InputField(prefix="The topic of the page: ", format=str) section = dspy.InputField(prefix="The section you need to write: ", format=str) output = dspy.OutputField( prefix="Write the section with proper inline citations (Start your writing with # section title. Don't include the page title or try to write other sections):\n", format=str, ) ================================================ FILE: knowledge_storm/storm_wiki/modules/article_polish.py ================================================ import copy from typing import Union import dspy from .storm_dataclass import StormArticle from ...interface import ArticlePolishingModule from ...utils import ArticleTextProcessing class StormArticlePolishingModule(ArticlePolishingModule): """ The interface for article generation stage. Given topic, collected information from knowledge curation stage, generated outline from outline generation stage. """ def __init__( self, article_gen_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], article_polish_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], ): self.article_gen_lm = article_gen_lm self.article_polish_lm = article_polish_lm self.polish_page = PolishPageModule( write_lead_engine=self.article_gen_lm, polish_engine=self.article_polish_lm ) def polish_article( self, topic: str, draft_article: StormArticle, remove_duplicate: bool = False ) -> StormArticle: """ Polish article. Args: topic (str): The topic of the article. draft_article (StormArticle): The draft article. remove_duplicate (bool): Whether to use one additional LM call to remove duplicates from the article. """ article_text = draft_article.to_string() polish_result = self.polish_page( topic=topic, draft_page=article_text, polish_whole_page=remove_duplicate ) lead_section = f"# summary\n{polish_result.lead_section}" polished_article = "\n\n".join([lead_section, polish_result.page]) polished_article_dict = ArticleTextProcessing.parse_article_into_dict( polished_article ) polished_article = copy.deepcopy(draft_article) polished_article.insert_or_create_section(article_dict=polished_article_dict) polished_article.post_processing() return polished_article class WriteLeadSection(dspy.Signature): """Write a lead section for the given Wikipedia page with the following guidelines: 1. The lead should stand on its own as a concise overview of the article's topic. It should identify the topic, establish context, explain why the topic is notable, and summarize the most important points, including any prominent controversies. 2. The lead section should be concise and contain no more than four well-composed paragraphs. 3. The lead section should be carefully sourced as appropriate. Add inline citations (e.g., "Washington, D.C., is the capital of the United States.[1][3].") where necessary. """ topic = dspy.InputField(prefix="The topic of the page: ", format=str) draft_page = dspy.InputField(prefix="The draft page:\n", format=str) lead_section = dspy.OutputField(prefix="Write the lead section:\n", format=str) class PolishPage(dspy.Signature): """You are a faithful text editor that is good at finding repeated information in the article and deleting them to make sure there is no repetition in the article. You won't delete any non-repeated part in the article. You will keep the inline citations and article structure (indicated by "#", "##", etc.) appropriately. Do your job for the following article.""" draft_page = dspy.InputField(prefix="The draft article:\n", format=str) page = dspy.OutputField(prefix="Your revised article:\n", format=str) class PolishPageModule(dspy.Module): def __init__( self, write_lead_engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], polish_engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], ): super().__init__() self.write_lead_engine = write_lead_engine self.polish_engine = polish_engine self.write_lead = dspy.Predict(WriteLeadSection) self.polish_page = dspy.Predict(PolishPage) def forward(self, topic: str, draft_page: str, polish_whole_page: bool = True): # NOTE: Change show_guidelines to false to make the generation more robust to different LM families. with dspy.settings.context(lm=self.write_lead_engine, show_guidelines=False): lead_section = self.write_lead( topic=topic, draft_page=draft_page ).lead_section if "The lead section:" in lead_section: lead_section = lead_section.split("The lead section:")[1].strip() if polish_whole_page: # NOTE: Change show_guidelines to false to make the generation more robust to different LM families. with dspy.settings.context(lm=self.polish_engine, show_guidelines=False): page = self.polish_page(draft_page=draft_page).page else: page = draft_page return dspy.Prediction(lead_section=lead_section, page=page) ================================================ FILE: knowledge_storm/storm_wiki/modules/callback.py ================================================ class BaseCallbackHandler: """Base callback handler that can be used to handle callbacks from the STORM pipeline.""" def on_identify_perspective_start(self, **kwargs): """Run when the perspective identification starts.""" pass def on_identify_perspective_end(self, perspectives: list[str], **kwargs): """Run when the perspective identification finishes.""" pass def on_information_gathering_start(self, **kwargs): """Run when the information gathering starts.""" pass def on_dialogue_turn_end(self, dlg_turn, **kwargs): """Run when a question asking and answering turn finishes.""" pass def on_information_gathering_end(self, **kwargs): """Run when the information gathering finishes.""" pass def on_information_organization_start(self, **kwargs): """Run when the information organization starts.""" pass def on_direct_outline_generation_end(self, outline: str, **kwargs): """Run when the direct outline generation finishes.""" pass def on_outline_refinement_end(self, outline: str, **kwargs): """Run when the outline refinement finishes.""" pass ================================================ FILE: knowledge_storm/storm_wiki/modules/knowledge_curation.py ================================================ import concurrent.futures import logging import os from concurrent.futures import as_completed from typing import Union, List, Tuple, Optional, Dict import dspy from .callback import BaseCallbackHandler from .persona_generator import StormPersonaGenerator from .storm_dataclass import DialogueTurn, StormInformationTable from ...interface import KnowledgeCurationModule, Retriever, Information from ...utils import ArticleTextProcessing try: from streamlit.runtime.scriptrunner import add_script_run_ctx streamlit_connection = True except ImportError as err: streamlit_connection = False script_dir = os.path.dirname(os.path.abspath(__file__)) class ConvSimulator(dspy.Module): """Simulate a conversation between a Wikipedia writer with specific persona and an expert.""" def __init__( self, topic_expert_engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], question_asker_engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], retriever: Retriever, max_search_queries_per_turn: int, search_top_k: int, max_turn: int, ): super().__init__() self.wiki_writer = WikiWriter(engine=question_asker_engine) self.topic_expert = TopicExpert( engine=topic_expert_engine, max_search_queries=max_search_queries_per_turn, search_top_k=search_top_k, retriever=retriever, ) self.max_turn = max_turn def forward( self, topic: str, persona: str, ground_truth_url: str, callback_handler: BaseCallbackHandler, ): """ topic: The topic to research. persona: The persona of the Wikipedia writer. ground_truth_url: The ground_truth_url will be excluded from search to avoid ground truth leakage in evaluation. """ dlg_history: List[DialogueTurn] = [] for _ in range(self.max_turn): user_utterance = self.wiki_writer( topic=topic, persona=persona, dialogue_turns=dlg_history ).question if user_utterance == "": logging.error("Simulated Wikipedia writer utterance is empty.") break if user_utterance.startswith("Thank you so much for your help!"): break expert_output = self.topic_expert( topic=topic, question=user_utterance, ground_truth_url=ground_truth_url ) dlg_turn = DialogueTurn( agent_utterance=expert_output.answer, user_utterance=user_utterance, search_queries=expert_output.queries, search_results=expert_output.searched_results, ) dlg_history.append(dlg_turn) callback_handler.on_dialogue_turn_end(dlg_turn=dlg_turn) return dspy.Prediction(dlg_history=dlg_history) class WikiWriter(dspy.Module): """Perspective-guided question asking in conversational setup. The asked question will be used to start a next round of information seeking.""" def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): super().__init__() self.ask_question_with_persona = dspy.ChainOfThought(AskQuestionWithPersona) self.ask_question = dspy.ChainOfThought(AskQuestion) self.engine = engine def forward( self, topic: str, persona: str, dialogue_turns: List[DialogueTurn], draft_page=None, ): conv = [] for turn in dialogue_turns[:-4]: conv.append( f"You: {turn.user_utterance}\nExpert: Omit the answer here due to space limit." ) for turn in dialogue_turns[-4:]: conv.append( f"You: {turn.user_utterance}\nExpert: {ArticleTextProcessing.remove_citations(turn.agent_utterance)}" ) conv = "\n".join(conv) conv = conv.strip() or "N/A" conv = ArticleTextProcessing.limit_word_count_preserve_newline(conv, 2500) with dspy.settings.context(lm=self.engine): if persona is not None and len(persona.strip()) > 0: question = self.ask_question_with_persona( topic=topic, persona=persona, conv=conv ).question else: question = self.ask_question( topic=topic, persona=persona, conv=conv ).question return dspy.Prediction(question=question) class AskQuestion(dspy.Signature): """You are an experienced Wikipedia writer. You are chatting with an expert to get information for the topic you want to contribute. Ask good questions to get more useful information relevant to the topic. When you have no more question to ask, say "Thank you so much for your help!" to end the conversation. Please only ask a question at a time and don't ask what you have asked before. Your questions should be related to the topic you want to write. """ topic = dspy.InputField(prefix="Topic you want to write: ", format=str) conv = dspy.InputField(prefix="Conversation history:\n", format=str) question = dspy.OutputField(format=str) class AskQuestionWithPersona(dspy.Signature): """You are an experienced Wikipedia writer and want to edit a specific page. Besides your identity as a Wikipedia writer, you have specific focus when researching the topic. Now, you are chatting with an expert to get information. Ask good questions to get more useful information. When you have no more question to ask, say "Thank you so much for your help!" to end the conversation. Please only ask a question at a time and don't ask what you have asked before. Your questions should be related to the topic you want to write. """ topic = dspy.InputField(prefix="Topic you want to write: ", format=str) persona = dspy.InputField( prefix="Your persona besides being a Wikipedia writer: ", format=str ) conv = dspy.InputField(prefix="Conversation history:\n", format=str) question = dspy.OutputField(format=str) class QuestionToQuery(dspy.Signature): """You want to answer the question using Google search. What do you type in the search box? Write the queries you will use in the following format: - query 1 - query 2 ... - query n""" topic = dspy.InputField(prefix="Topic you are discussing about: ", format=str) question = dspy.InputField(prefix="Question you want to answer: ", format=str) queries = dspy.OutputField(format=str) class AnswerQuestion(dspy.Signature): """You are an expert who can use information effectively. You are chatting with a Wikipedia writer who wants to write a Wikipedia page on topic you know. You have gathered the related information and will now use the information to form a response. Make your response as informative as possible, ensuring that every sentence is supported by the gathered information. If the [gathered information] is not directly related to the [topic] or [question], provide the most relevant answer based on the available information. If no appropriate answer can be formulated, respond with, “I cannot answer this question based on the available information,” and explain any limitations or gaps. """ topic = dspy.InputField(prefix="Topic you are discussing about:", format=str) conv = dspy.InputField(prefix="Question:\n", format=str) info = dspy.InputField(prefix="Gathered information:\n", format=str) answer = dspy.OutputField( prefix="Now give your response. (Try to use as many different sources as possible and add do not hallucinate.)\n", format=str, ) class TopicExpert(dspy.Module): """Answer questions using search-based retrieval and answer generation. This module conducts the following steps: 1. Generate queries from the question. 2. Search for information using the queries. 3. Filter out unreliable sources. 4. Generate an answer using the retrieved information. """ def __init__( self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel], max_search_queries: int, search_top_k: int, retriever: Retriever, ): super().__init__() self.generate_queries = dspy.Predict(QuestionToQuery) self.retriever = retriever self.answer_question = dspy.Predict(AnswerQuestion) self.engine = engine self.max_search_queries = max_search_queries self.search_top_k = search_top_k def forward(self, topic: str, question: str, ground_truth_url: str): with dspy.settings.context(lm=self.engine, show_guidelines=False): # Identify: Break down question into queries. queries = self.generate_queries(topic=topic, question=question).queries queries = [ q.replace("-", "").strip().strip('"').strip('"').strip() for q in queries.split("\n") ] queries = queries[: self.max_search_queries] # Search searched_results: List[Information] = self.retriever.retrieve( list(set(queries)), exclude_urls=[ground_truth_url] ) if len(searched_results) > 0: # Evaluate: Simplify this part by directly using the top 1 snippet. info = "" for n, r in enumerate(searched_results): info += "\n".join(f"[{n + 1}]: {s}" for s in r.snippets[:1]) info += "\n\n" info = ArticleTextProcessing.limit_word_count_preserve_newline( info, 1000 ) try: answer = self.answer_question( topic=topic, conv=question, info=info ).answer answer = ArticleTextProcessing.remove_uncompleted_sentences_with_citations( answer ) except Exception as e: logging.error(f"Error occurs when generating answer: {e}") answer = "Sorry, I cannot answer this question. Please ask another question." else: # When no information is found, the expert shouldn't hallucinate. answer = "Sorry, I cannot find information for this question. Please ask another question." return dspy.Prediction( queries=queries, searched_results=searched_results, answer=answer ) class StormKnowledgeCurationModule(KnowledgeCurationModule): """ The interface for knowledge curation stage. Given topic, return collected information. """ def __init__( self, retriever: Retriever, persona_generator: Optional[StormPersonaGenerator], conv_simulator_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], question_asker_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel], max_search_queries_per_turn: int, search_top_k: int, max_conv_turn: int, max_thread_num: int, ): """ Store args and finish initialization. """ self.retriever = retriever self.persona_generator = persona_generator self.conv_simulator_lm = conv_simulator_lm self.search_top_k = search_top_k self.max_thread_num = max_thread_num self.retriever = retriever self.conv_simulator = ConvSimulator( topic_expert_engine=conv_simulator_lm, question_asker_engine=question_asker_lm, retriever=retriever, max_search_queries_per_turn=max_search_queries_per_turn, search_top_k=search_top_k, max_turn=max_conv_turn, ) def _get_considered_personas(self, topic: str, max_num_persona) -> List[str]: return self.persona_generator.generate_persona( topic=topic, max_num_persona=max_num_persona ) def _run_conversation( self, conv_simulator, topic, ground_truth_url, considered_personas, callback_handler: BaseCallbackHandler, ) -> List[Tuple[str, List[DialogueTurn]]]: """ Executes multiple conversation simulations concurrently, each with a different persona, and collects their dialog histories. The dialog history of each conversation is cleaned up before being stored. Parameters: conv_simulator (callable): The function to simulate conversations. It must accept four parameters: `topic`, `ground_truth_url`, `persona`, and `callback_handler`, and return an object that has a `dlg_history` attribute. topic (str): The topic of conversation for the simulations. ground_truth_url (str): The URL to the ground truth data related to the conversation topic. considered_personas (list): A list of personas under which the conversation simulations will be conducted. Each persona is passed to `conv_simulator` individually. callback_handler (callable): A callback function that is passed to `conv_simulator`. It should handle any callbacks or events during the simulation. Returns: list of tuples: A list where each tuple contains a persona and its corresponding cleaned dialog history (`dlg_history`) from the conversation simulation. """ conversations = [] def run_conv(persona): return conv_simulator( topic=topic, ground_truth_url=ground_truth_url, persona=persona, callback_handler=callback_handler, ) max_workers = min(self.max_thread_num, len(considered_personas)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: future_to_persona = { executor.submit(run_conv, persona): persona for persona in considered_personas } if streamlit_connection: # Ensure the logging context is correct when connecting with Streamlit frontend. for t in executor._threads: add_script_run_ctx(t) for future in as_completed(future_to_persona): persona = future_to_persona[future] conv = future.result() conversations.append( (persona, ArticleTextProcessing.clean_up_citation(conv).dlg_history) ) return conversations def research( self, topic: str, ground_truth_url: str, callback_handler: BaseCallbackHandler, max_perspective: int = 0, disable_perspective: bool = True, return_conversation_log=False, ) -> Union[StormInformationTable, Tuple[StormInformationTable, Dict]]: """ Curate information and knowledge for the given topic Args: topic: topic of interest in natural language. Returns: collected_information: collected information in InformationTable type. """ # identify personas callback_handler.on_identify_perspective_start() considered_personas = [] if disable_perspective: considered_personas = [""] else: considered_personas = self._get_considered_personas( topic=topic, max_num_persona=max_perspective ) callback_handler.on_identify_perspective_end(perspectives=considered_personas) # run conversation callback_handler.on_information_gathering_start() conversations = self._run_conversation( conv_simulator=self.conv_simulator, topic=topic, ground_truth_url=ground_truth_url, considered_personas=considered_personas, callback_handler=callback_handler, ) information_table = StormInformationTable(conversations) callback_handler.on_information_gathering_end() if return_conversation_log: return information_table, StormInformationTable.construct_log_dict( conversations ) return information_table ================================================ FILE: knowledge_storm/storm_wiki/modules/outline_generation.py ================================================ from typing import Union, Optional, Tuple import dspy from .callback import BaseCallbackHandler from .storm_dataclass import StormInformationTable, StormArticle from ...interface import OutlineGenerationModule from ...utils import ArticleTextProcessing class StormOutlineGenerationModule(OutlineGenerationModule): """ The interface for outline generation stage. Given topic, collected information from knowledge curation stage, generate outline for the article. """ def __init__(self, outline_gen_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel]): super().__init__() self.outline_gen_lm = outline_gen_lm self.write_outline = WriteOutline(engine=self.outline_gen_lm) def generate_outline( self, topic: str, information_table: StormInformationTable, old_outline: Optional[StormArticle] = None, callback_handler: BaseCallbackHandler = None, return_draft_outline=False, ) -> Union[StormArticle, Tuple[StormArticle, StormArticle]]: """ Generates an outline for an article based on the specified topic and the information gathered during the knowledge curation stage. This method can optionally return both the final article outline and a draft outline if required. Args: topic (str): The topic of the article. information_table (StormInformationTable): The information table containing the collected information. old_outline (Optional[StormArticle]): An optional previous version of the article outline that can be used for reference or comparison. Defaults to None. callback_handler (BaseCallbackHandler): An optional callback handler that can be used to trigger custom callbacks at various stages of the outline generation process, such as when the information organization starts. Defaults to None. return_draft_outline (bool): A flag indicating whether the method should return both the final article outline and a draft version of the outline. If False, only the final article outline is returned. Defaults to False. Returns: Union[StormArticle, Tuple[StormArticle, StormArticle]]: Depending on the value of `return_draft_outline`, this method returns either a single `StormArticle` object containing the final outline or a tuple of two `StormArticle` objects, the first containing the final outline and the second containing the draft outline. """ if callback_handler is not None: callback_handler.on_information_organization_start() concatenated_dialogue_turns = sum( [conv for (_, conv) in information_table.conversations], [] ) result = self.write_outline( topic=topic, dlg_history=concatenated_dialogue_turns, callback_handler=callback_handler, ) article_with_outline_only = StormArticle.from_outline_str( topic=topic, outline_str=result.outline ) article_with_draft_outline_only = StormArticle.from_outline_str( topic=topic, outline_str=result.old_outline ) if not return_draft_outline: return article_with_outline_only return article_with_outline_only, article_with_draft_outline_only class WriteOutline(dspy.Module): """Generate the outline for the Wikipedia page.""" def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): super().__init__() self.draft_page_outline = dspy.Predict(WritePageOutline) self.write_page_outline = dspy.Predict(WritePageOutlineFromConv) self.engine = engine def forward( self, topic: str, dlg_history, old_outline: Optional[str] = None, callback_handler: BaseCallbackHandler = None, ): trimmed_dlg_history = [] for turn in dlg_history: if ( "topic you" in turn.agent_utterance.lower() or "topic you" in turn.user_utterance.lower() ): continue trimmed_dlg_history.append(turn) conv = "\n".join( [ f"Wikipedia Writer: {turn.user_utterance}\nExpert: {turn.agent_utterance}" for turn in trimmed_dlg_history ] ) conv = ArticleTextProcessing.remove_citations(conv) conv = ArticleTextProcessing.limit_word_count_preserve_newline(conv, 5000) with dspy.settings.context(lm=self.engine): if old_outline is None: old_outline = ArticleTextProcessing.clean_up_outline( self.draft_page_outline(topic=topic).outline ) if callback_handler: callback_handler.on_direct_outline_generation_end( outline=old_outline ) outline = ArticleTextProcessing.clean_up_outline( self.write_page_outline( topic=topic, old_outline=old_outline, conv=conv ).outline ) if callback_handler: callback_handler.on_outline_refinement_end(outline=outline) return dspy.Prediction(outline=outline, old_outline=old_outline) class WritePageOutline(dspy.Signature): """Write an outline for a Wikipedia page. Here is the format of your writing: 1. Use "#" Title" to indicate section title, "##" Title" to indicate subsection title, "###" Title" to indicate subsubsection title, and so on. 2. Do not include other information. 3. Do not include topic name itself in the outline. """ topic = dspy.InputField(prefix="The topic you want to write: ", format=str) outline = dspy.OutputField(prefix="Write the Wikipedia page outline:\n", format=str) class NaiveOutlineGen(dspy.Module): """Generate the outline with LLM's parametric knowledge directly.""" def __init__(self): super().__init__() self.write_outline = dspy.Predict(WritePageOutline) def forward(self, topic: str): outline = self.write_outline(topic=topic).outline return dspy.Prediction(outline=outline) class WritePageOutlineFromConv(dspy.Signature): """Improve an outline for a Wikipedia page. You already have a draft outline that covers the general information. Now you want to improve it based on the information learned from an information-seeking conversation to make it more informative. Here is the format of your writing: 1. Use "#" Title" to indicate section title, "##" Title" to indicate subsection title, "###" Title" to indicate subsubsection title, and so on. 2. Do not include other information. 3. Do not include topic name itself in the outline. """ topic = dspy.InputField(prefix="The topic you want to write: ", format=str) conv = dspy.InputField(prefix="Conversation history:\n", format=str) old_outline = dspy.OutputField(prefix="Current outline:\n", format=str) outline = dspy.OutputField( prefix='Write the Wikipedia page outline (Use "#" Title" to indicate section title, "##" Title" to indicate subsection title, ...):\n', format=str, ) ================================================ FILE: knowledge_storm/storm_wiki/modules/persona_generator.py ================================================ import logging import re from typing import Union, List import dspy import requests from bs4 import BeautifulSoup def get_wiki_page_title_and_toc(url): """Get the main title and table of contents from an url of a Wikipedia page.""" response = requests.get(url) soup = BeautifulSoup(response.content, "html.parser") # Get the main title from the first h1 tag main_title = soup.find("h1").text.replace("[edit]", "").strip().replace("\xa0", " ") toc = "" levels = [] excluded_sections = { "Contents", "See also", "Notes", "References", "External links", } # Start processing from h2 to exclude the main title from TOC for header in soup.find_all(["h2", "h3", "h4", "h5", "h6"]): level = int( header.name[1] ) # Extract the numeric part of the header tag (e.g., '2' from 'h2') section_title = header.text.replace("[edit]", "").strip().replace("\xa0", " ") if section_title in excluded_sections: continue while levels and level <= levels[-1]: levels.pop() levels.append(level) indentation = " " * (len(levels) - 1) toc += f"{indentation}{section_title}\n" return main_title, toc.strip() class FindRelatedTopic(dspy.Signature): """I'm writing a Wikipedia page for a topic mentioned below. Please identify and recommend some Wikipedia pages on closely related subjects. I'm looking for examples that provide insights into interesting aspects commonly associated with this topic, or examples that help me understand the typical content and structure included in Wikipedia pages for similar topics. Please list the urls in separate lines.""" topic = dspy.InputField(prefix="Topic of interest:", format=str) related_topics = dspy.OutputField(format=str) class GenPersona(dspy.Signature): """You need to select a group of Wikipedia editors who will work together to create a comprehensive article on the topic. Each of them represents a different perspective, role, or affiliation related to this topic. You can use other Wikipedia pages of related topics for inspiration. For each editor, add a description of what they will focus on. Give your answer in the following format: 1. short summary of editor 1: description\n2. short summary of editor 2: description\n... """ topic = dspy.InputField(prefix="Topic of interest:", format=str) examples = dspy.InputField( prefix="Wiki page outlines of related topics for inspiration:\n", format=str ) personas = dspy.OutputField(format=str) class CreateWriterWithPersona(dspy.Module): """Discover different perspectives of researching the topic by reading Wikipedia pages of related topics.""" def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): super().__init__() self.find_related_topic = dspy.ChainOfThought(FindRelatedTopic) self.gen_persona = dspy.ChainOfThought(GenPersona) self.engine = engine def forward(self, topic: str, draft=None): with dspy.settings.context(lm=self.engine): # Get section names from wiki pages of relevant topics for inspiration. related_topics = self.find_related_topic(topic=topic).related_topics urls = [] for s in related_topics.split("\n"): if "http" in s: urls.append(s[s.find("http") :]) examples = [] for url in urls: try: title, toc = get_wiki_page_title_and_toc(url) examples.append(f"Title: {title}\nTable of Contents: {toc}") except Exception as e: logging.error(f"Error occurs when processing {url}: {e}") continue if len(examples) == 0: examples.append("N/A") gen_persona_output = self.gen_persona( topic=topic, examples="\n----------\n".join(examples) ).personas personas = [] for s in gen_persona_output.split("\n"): match = re.search(r"\d+\.\s*(.*)", s) if match: personas.append(match.group(1)) sorted_personas = personas return dspy.Prediction( personas=personas, raw_personas_output=sorted_personas, related_topics=related_topics, ) class StormPersonaGenerator: """ A generator class for creating personas based on a given topic. This class uses an underlying engine to generate personas tailored to the specified topic. The generator integrates with a `CreateWriterWithPersona` instance to create diverse personas, including a default 'Basic fact writer' persona. Attributes: create_writer_with_persona (CreateWriterWithPersona): An instance responsible for generating personas based on the provided engine and topic. Args: engine (Union[dspy.dsp.LM, dspy.dsp.HFModel]): The underlying engine used for generating personas. It must be an instance of either `dspy.dsp.LM` or `dspy.dsp.HFModel`. """ def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]): self.create_writer_with_persona = CreateWriterWithPersona(engine=engine) def generate_persona(self, topic: str, max_num_persona: int = 3) -> List[str]: """ Generates a list of personas based on the provided topic, up to a maximum number specified. This method first creates personas using the underlying `create_writer_with_persona` instance and then prepends a default 'Basic fact writer' persona to the list before returning it. The number of personas returned is limited to `max_num_persona`, excluding the default persona. Args: topic (str): The topic for which personas are to be generated. max_num_persona (int): The maximum number of personas to generate, excluding the default 'Basic fact writer' persona. Returns: List[str]: A list of persona descriptions, including the default 'Basic fact writer' persona and up to `max_num_persona` additional personas generated based on the topic. """ personas = self.create_writer_with_persona(topic=topic) default_persona = "Basic fact writer: Basic fact writer focusing on broadly covering the basic facts about the topic." considered_personas = [default_persona] + personas.personas[:max_num_persona] return considered_personas ================================================ FILE: knowledge_storm/storm_wiki/modules/retriever.py ================================================ from typing import Union, List from urllib.parse import urlparse import dspy from ...interface import Retriever, Information from ...utils import ArticleTextProcessing # Internet source restrictions according to Wikipedia standard: # https://en.wikipedia.org/wiki/Wikipedia:Reliable_sources/Perennial_sources GENERALLY_UNRELIABLE = { "112_Ukraine", "Ad_Fontes_Media", "AlterNet", "Amazon", "Anadolu_Agency_(controversial_topics)", "Ancestry.com", "Answers.com", "Antiwar.com", "Anti-Defamation_League", "arXiv", "Atlas_Obscura_places", "Bild", "Blaze_Media", "Blogger", "BroadwayWorld", "California_Globe", "The_Canary", "CelebrityNetWorth", "CESNUR", "ChatGPT", "CNET_(November_2022\u2013present)", "CoinDesk", "Consortium_News", "CounterPunch", "Correo_del_Orinoco", "Cracked.com", "Daily_Express", "Daily_Kos", "Daily_Sabah", "The_Daily_Wire", "Discogs", "Distractify", "The_Electronic_Intifada", "Encyclopaedia_Metallum", "Ethnicity_of_Celebs", "Facebook", "FamilySearch", "Fandom", "The_Federalist", "Find_a_Grave", "Findmypast", "Flags_of_the_World", "Flickr", "Forbes.com_contributors", "Fox_News_(politics_and_science)", "Fox_News_(talk_shows)", "Gawker", "GB_News", "Geni.com", "gnis-class", "gns-class", "GlobalSecurity.org", "Goodreads", "Guido_Fawkes", "Heat_Street", "History", "HuffPost_contributors", "IMDb", "Independent_Media_Center", "Inquisitr", "International_Business_Times", "Investopedia", "Jewish_Virtual_Library", "Joshua_Project", "Know_Your_Meme", "Land_Transport_Guru", "LinkedIn", "LiveJournal", "Marquis_Who's_Who", "Mashable_sponsored_content", "MEAWW", "Media_Bias/Fact_Check", "Media_Research_Center", "Medium", "metal-experience", "Metro", "The_New_American", "New_York_Post", "NGO_Monitor", "The_Onion", "Our_Campaigns", "PanAm_Post", "Patheos", "An_Phoblacht", "The_Post_Millennial", "arXiv", "bioRxiv", "medRxiv", "PeerJ Preprints", "Preprints.org", "SSRN", "PR_Newswire", "Quadrant", "Quillette", "Quora", "Raw_Story", "Reddit", "RedState", "ResearchGate", "Rolling_Stone_(politics_and_society,_2011\u2013present)", "Rolling_Stone_(Culture_Council)", "Scribd", "Scriptural_texts", "Simple_Flying", "Sixth_Tone_(politics)", "The_Skwawkbox", "SourceWatch", "Spirit_of_Metal", "Sportskeeda", "Stack_Exchange", "Stack_Overflow", "MathOverflow", "Ask_Ubuntu", "starsunfolded.com", "Statista", "TASS", "The_Truth_About_Guns", "TV.com", "TV_Tropes", "Twitter", "X.com", "Urban_Dictionary", "Venezuelanalysis", "VGChartz", "VoC", "Washington_Free_Beacon", "Weather2Travel", "The_Western_Journal", "We_Got_This_Covered", "WhatCulture", "Who's_Who_(UK)", "WhoSampled", "Wikidata", "WikiLeaks", "Wikinews", "Wikipedia", "WordPress.com", "Worldometer", "YouTube", "ZDNet", } DEPRECATED = { "Al_Mayadeen", "ANNA_News", "Baidu_Baike", "China_Global_Television_Network", "The_Cradle", "Crunchbase", "The_Daily_Caller", "Daily_Mail", "Daily_Star", "The_Epoch_Times", "FrontPage_Magazine", "The_Gateway_Pundit", "Global_Times", "The_Grayzone", "HispanTV", "Jihad_Watch", "Last.fm", "LifeSiteNews", "The_Mail_on_Sunday", "MintPress_News", "National_Enquirer", "New_Eastern_Outlook", "News_Break", "NewsBlaze", "News_of_the_World", "Newsmax", "NNDB", "Occupy_Democrats", "Office_of_Cuba_Broadcasting", "One_America_News_Network", "Peerage_websites", "Press_TV", "Project_Veritas", "Rate_Your_Music", "Republic_TV", "Royal_Central", "RT", "Sputnik", "The_Sun", "Taki's_Magazine", "Tasnim_News_Agency", "Telesur", "The_Unz_Review", "VDARE", "Voltaire_Network", "WorldNetDaily", "Zero_Hedge", } BLACKLISTED = { "Advameg", "bestgore.com", "Breitbart_News", "Centre_for_Research_on_Globalization", "Examiner.com", "Famous_Birthdays", "Healthline", "InfoWars", "Lenta.ru", "LiveLeak", "Lulu.com", "MyLife", "Natural_News", "OpIndia", "The_Points_Guy", "The_Points_Guy_(sponsored_content)", "Swarajya", "Veterans_Today", "ZoomInfo", } def is_valid_wikipedia_source(url): parsed_url = urlparse(url) # Check if the URL is from a reliable domain combined_set = GENERALLY_UNRELIABLE | DEPRECATED | BLACKLISTED for domain in combined_set: if domain in parsed_url.netloc: return False return True ================================================ FILE: knowledge_storm/storm_wiki/modules/storm_dataclass.py ================================================ import copy import re from collections import OrderedDict from typing import Union, Optional, Any, List, Tuple, Dict import numpy as np from sentence_transformers import SentenceTransformer from sklearn.metrics.pairwise import cosine_similarity from ...interface import Information, InformationTable, Article, ArticleSectionNode from ...utils import ArticleTextProcessing, FileIOHelper class DialogueTurn: def __init__( self, agent_utterance: str = None, user_utterance: str = None, search_queries: Optional[List[str]] = None, search_results: Optional[List[Union[Information, Dict]]] = None, ): self.agent_utterance = agent_utterance self.user_utterance = user_utterance self.search_queries = search_queries self.search_results = search_results if self.search_results: for idx in range(len(self.search_results)): if type(self.search_results[idx]) == dict: self.search_results[idx] = Information.from_dict( self.search_results[idx] ) def log(self): """ Returns a json object that contains all information inside `self` """ return OrderedDict( { "agent_utterance": self.agent_utterance, "user_utterance": self.user_utterance, "search_queries": self.search_queries, "search_results": [data.to_dict() for data in self.search_results], } ) class StormInformationTable(InformationTable): """ The InformationTable class serves as data class to store the information collected during KnowledgeCuration stage. Create subclass to incorporate more information as needed. For example, in STORM paper https://arxiv.org/pdf/2402.14207.pdf, additional information would be perspective guided dialogue history. """ def __init__(self, conversations=List[Tuple[str, List[DialogueTurn]]]): super().__init__() self.conversations = conversations self.url_to_info: Dict[str, Information] = ( StormInformationTable.construct_url_to_info(self.conversations) ) @staticmethod def construct_url_to_info( conversations: List[Tuple[str, List[DialogueTurn]]] ) -> Dict[str, Information]: url_to_info = {} for persona, conv in conversations: for turn in conv: for storm_info in turn.search_results: if storm_info.url in url_to_info: url_to_info[storm_info.url].snippets.extend(storm_info.snippets) else: url_to_info[storm_info.url] = storm_info for url in url_to_info: url_to_info[url].snippets = list(set(url_to_info[url].snippets)) return url_to_info @staticmethod def construct_log_dict( conversations: List[Tuple[str, List[DialogueTurn]]] ) -> List[Dict[str, Union[str, Any]]]: conversation_log = [] for persona, conv in conversations: conversation_log.append( {"perspective": persona, "dlg_turns": [turn.log() for turn in conv]} ) return conversation_log def dump_url_to_info(self, path): url_to_info = copy.deepcopy(self.url_to_info) for url in url_to_info: url_to_info[url] = url_to_info[url].to_dict() FileIOHelper.dump_json(url_to_info, path) @classmethod def from_conversation_log_file(cls, path): conversation_log_data = FileIOHelper.load_json(path) conversations = [] for item in conversation_log_data: dialogue_turns = [DialogueTurn(**turn) for turn in item["dlg_turns"]] persona = item["perspective"] conversations.append((persona, dialogue_turns)) return cls(conversations) def prepare_table_for_retrieval(self): self.encoder = SentenceTransformer("paraphrase-MiniLM-L6-v2") self.collected_urls = [] self.collected_snippets = [] for url, information in self.url_to_info.items(): for snippet in information.snippets: self.collected_urls.append(url) self.collected_snippets.append(snippet) self.encoded_snippets = self.encoder.encode(self.collected_snippets) def retrieve_information( self, queries: Union[List[str], str], search_top_k ) -> List[Information]: selected_urls = [] selected_snippets = [] if type(queries) is str: queries = [queries] for query in queries: encoded_query = self.encoder.encode(query) sim = cosine_similarity([encoded_query], self.encoded_snippets)[0] sorted_indices = np.argsort(sim) for i in sorted_indices[-search_top_k:][::-1]: selected_urls.append(self.collected_urls[i]) selected_snippets.append(self.collected_snippets[i]) url_to_snippets = {} for url, snippet in zip(selected_urls, selected_snippets): if url not in url_to_snippets: url_to_snippets[url] = set() url_to_snippets[url].add(snippet) selected_url_to_info = {} for url in url_to_snippets: selected_url_to_info[url] = copy.deepcopy(self.url_to_info[url]) selected_url_to_info[url].snippets = list(url_to_snippets[url]) return list(selected_url_to_info.values()) class StormArticle(Article): def __init__(self, topic_name): super().__init__(topic_name=topic_name) self.reference = {"url_to_unified_index": {}, "url_to_info": {}} def find_section( self, node: ArticleSectionNode, name: str ) -> Optional[ArticleSectionNode]: """ Return the node of the section given the section name. Args: node: the node as the root to find. name: the name of node as section name Return: reference of the node or None if section name has no match """ if node.section_name == name: return node for child in node.children: result = self.find_section(child, name) if result: return result return None def _merge_new_info_to_references( self, new_info_list: List[Information], index_to_keep=None ) -> Dict[int, int]: """ Merges new storm information into existing references and updates the citation index mapping. Args: new_info_list (List[Information]): A list of dictionaries representing new storm information. index_to_keep (List[int]): A list of index of the new_info_list to keep. If none, keep all. Returns: Dict[int, int]: A dictionary mapping the index of each storm information piece in the input list to its unified citation index in the references. """ citation_idx_mapping = {} for idx, storm_info in enumerate(new_info_list): if index_to_keep is not None and idx not in index_to_keep: continue url = storm_info.url if url not in self.reference["url_to_unified_index"]: self.reference["url_to_unified_index"][url] = ( len(self.reference["url_to_unified_index"]) + 1 ) # The citation index starts from 1. self.reference["url_to_info"][url] = storm_info else: existing_snippets = self.reference["url_to_info"][url].snippets existing_snippets.extend(storm_info.snippets) self.reference["url_to_info"][url].snippets = list( set(existing_snippets) ) citation_idx_mapping[idx + 1] = self.reference["url_to_unified_index"][ url ] # The citation index starts from 1. return citation_idx_mapping def insert_or_create_section( self, article_dict: Dict[str, Dict], parent_section_name: str = None, trim_children=False, ): parent_node = ( self.root if parent_section_name is None else self.find_section(self.root, parent_section_name) ) if trim_children: section_names = set(article_dict.keys()) for child in parent_node.children[:]: if child.section_name not in section_names: parent_node.remove_child(child) for section_name, content_dict in article_dict.items(): current_section_node = self.find_section(parent_node, section_name) if current_section_node is None: current_section_node = ArticleSectionNode( section_name=section_name, content=content_dict["content"].strip() ) insert_to_front = ( parent_node.section_name == self.root.section_name and current_section_node.section_name == "summary" ) parent_node.add_child( current_section_node, insert_to_front=insert_to_front ) else: current_section_node.content = content_dict["content"].strip() self.insert_or_create_section( article_dict=content_dict["subsections"], parent_section_name=section_name, trim_children=True, ) def update_section( self, current_section_content: str, current_section_info_list: List[Information], parent_section_name: Optional[str] = None, ) -> Optional[ArticleSectionNode]: """ Add new section to the article. Args: current_section_name: new section heading name in string format. parent_section_name: under which parent section to add the new one. Default to root. current_section_content: optional section content. Returns: the ArticleSectionNode for current section if successfully created / updated. Otherwise none. """ if current_section_info_list is not None: references = set( [int(x) for x in re.findall(r"\[(\d+)\]", current_section_content)] ) # for any reference number greater than max number of references, delete the reference if len(references) > 0: max_ref_num = max(references) if max_ref_num > len(current_section_info_list): for i in range(len(current_section_info_list), max_ref_num + 1): current_section_content = current_section_content.replace( f"[{i}]", "" ) if i in references: references.remove(i) # for any reference that is not used, trim it from current_section_info_list index_to_keep = [i - 1 for i in references] citation_mapping = self._merge_new_info_to_references( current_section_info_list, index_to_keep ) current_section_content = ArticleTextProcessing.update_citation_index( current_section_content, citation_mapping ) if parent_section_name is None: parent_section_name = self.root.section_name article_dict = ArticleTextProcessing.parse_article_into_dict( current_section_content ) self.insert_or_create_section( article_dict=article_dict, parent_section_name=parent_section_name, trim_children=False, ) def get_outline_as_list( self, root_section_name: Optional[str] = None, add_hashtags: bool = False, include_root: bool = True, ) -> List[str]: """ Get outline of the article as a list. Args: section_name: get all section names in pre-order travel ordering in the subtree of section_name. For example: #root ##section1 ###section1.1 ###section1.2 ##section2 article.get_outline_as_list("section1") returns [section1, section1.1, section1.2, section2] Returns: list of section and subsection names. """ if root_section_name is None: section_node = self.root else: section_node = self.find_section(self.root, root_section_name) include_root = include_root or section_node != self.root.section_name if section_node is None: return [] result = [] def preorder_traverse(node, level): prefix = ( "#" * level if add_hashtags else "" ) # Adjust level if excluding root result.append( f"{prefix} {node.section_name}".strip() if add_hashtags else node.section_name ) for child in node.children: preorder_traverse(child, level + 1) # Adjust the initial level based on whether root is included and hashtags are added if include_root: preorder_traverse(section_node, level=1) else: for child in section_node.children: preorder_traverse(child, level=1) return result def to_string(self) -> str: """ Get outline of the article as a list. Returns: list of section and subsection names. """ result = [] def preorder_traverse(node, level): prefix = "#" * level result.append(f"{prefix} {node.section_name}".strip()) result.append(node.content) for child in node.children: preorder_traverse(child, level + 1) # Adjust the initial level based on whether root is included and hashtags are added for child in self.root.children: preorder_traverse(child, level=1) result = [i.strip() for i in result if i is not None and i.strip()] return "\n\n".join(result) def reorder_reference_index(self): # pre-order traversal to get order of references appear in the article ref_indices = [] def pre_order_find_index(node): if node is not None: if node.content is not None and node.content: ref_indices.extend( ArticleTextProcessing.parse_citation_indices(node.content) ) for child in node.children: pre_order_find_index(child) pre_order_find_index(self.root) # constrcut index mapping ref_index_mapping = {} for ref_index in ref_indices: if ref_index not in ref_index_mapping: ref_index_mapping[ref_index] = len(ref_index_mapping) + 1 # update content def pre_order_update_index(node): if node is not None: if node.content is not None and node.content: node.content = ArticleTextProcessing.update_citation_index( node.content, ref_index_mapping ) for child in node.children: pre_order_update_index(child) pre_order_update_index(self.root) # update reference for url in list(self.reference["url_to_unified_index"]): pre_index = self.reference["url_to_unified_index"][url] if pre_index not in ref_index_mapping: del self.reference["url_to_unified_index"][url] else: new_index = ref_index_mapping[pre_index] self.reference["url_to_unified_index"][url] = new_index def get_outline_tree(self): def build_tree(node) -> Dict[str, Dict]: tree = {} for child in node.children: tree[child.section_name] = build_tree(child) return tree if tree else {} return build_tree(self.root) def get_first_level_section_names(self) -> List[str]: """ Get first level section names """ return [i.section_name for i in self.root.children] @classmethod def from_outline_file(cls, topic: str, file_path: str): """ Create StormArticle class instance from outline file. """ outline_str = FileIOHelper.load_str(file_path) return StormArticle.from_outline_str(topic=topic, outline_str=outline_str) @classmethod def from_outline_str(cls, topic: str, outline_str: str): """ Create StormArticle class instance from outline only string. """ lines = [] try: lines = outline_str.split("\n") lines = [line.strip() for line in lines if line.strip()] except: pass instance = cls(topic) if lines: a = lines[0].startswith("#") and lines[0].replace("#", "").strip().lower() b = topic.lower().replace("_", " ") adjust_level = lines[0].startswith("#") and lines[0].replace( "#", "" ).strip().lower() == topic.lower().replace("_", " ") if adjust_level: lines = lines[1:] node_stack = [(0, instance.root)] # Stack to keep track of (level, node) for line in lines: level = line.count("#") - adjust_level section_name = line.replace("#", "").strip() if section_name == topic: continue new_node = ArticleSectionNode(section_name) while node_stack and level <= node_stack[-1][0]: node_stack.pop() node_stack[-1][1].add_child(new_node) node_stack.append((level, new_node)) return instance def dump_outline_to_file(self, file_path): outline = self.get_outline_as_list(add_hashtags=True, include_root=False) FileIOHelper.write_str("\n".join(outline), file_path) def dump_reference_to_file(self, file_path): reference = copy.deepcopy(self.reference) for url in reference["url_to_info"]: reference["url_to_info"][url] = reference["url_to_info"][url].to_dict() FileIOHelper.dump_json(reference, file_path) def dump_article_as_plain_text(self, file_path): text = self.to_string() FileIOHelper.write_str(text, file_path) @classmethod def from_string(cls, topic_name: str, article_text: str, references: dict): article_dict = ArticleTextProcessing.parse_article_into_dict(article_text) article = cls(topic_name=topic_name) article.insert_or_create_section(article_dict=article_dict) for url in list(references["url_to_info"]): references["url_to_info"][url] = Information.from_dict( references["url_to_info"][url] ) article.reference = references return article def post_processing(self): self.prune_empty_nodes() self.reorder_reference_index() ================================================ FILE: knowledge_storm/utils.py ================================================ import concurrent.futures import dspy import httpx import json import logging import os import pickle import re import regex import sys import toml from typing import List, Dict from tqdm import tqdm from langchain_text_splitters import RecursiveCharacterTextSplitter from trafilatura import extract from .lm import LitellmModel logging.getLogger("httpx").setLevel(logging.WARNING) # Disable INFO logging for httpx. def truncate_filename(filename, max_length=125): """Truncate filename to max_length to ensure the filename won't exceed the file system limit. Args: filename: str max_length: int, default to 125 (usual path length limit is 255 chars) """ if len(filename) > max_length: truncated_filename = filename[:max_length] logging.warning( f"Filename is too long. Filename is truncated to {truncated_filename}." ) return truncated_filename return filename def load_api_key(toml_file_path): try: with open(toml_file_path, "r") as file: data = toml.load(file) except FileNotFoundError: print(f"File not found: {toml_file_path}", file=sys.stderr) return except toml.TomlDecodeError: print(f"Error decoding TOML file: {toml_file_path}", file=sys.stderr) return # Set environment variables for key, value in data.items(): os.environ[key] = str(value) def makeStringRed(message): return f"\033[91m {message}\033[00m" class QdrantVectorStoreManager: """ Helper class for managing the Qdrant vector store, can be used with `VectorRM` in rm.py. Before you initialize `VectorRM`, call `create_or_update_vector_store` to create or update the vector store. Once you have the vector store, you can initialize `VectorRM` with the vector store path or the Qdrant server URL. """ @staticmethod def _check_create_collection( client: "QdrantClient", collection_name: str, model: "HuggingFaceEmbeddings" ): from langchain_qdrant import Qdrant from qdrant_client import models """Check if the Qdrant collection exists and create it if it does not.""" if client is None: raise ValueError("Qdrant client is not initialized.") if client.collection_exists(collection_name=f"{collection_name}"): print(f"Collection {collection_name} exists. Loading the collection...") return Qdrant( client=client, collection_name=collection_name, embeddings=model, ) else: print( f"Collection {collection_name} does not exist. Creating the collection..." ) # create the collection client.create_collection( collection_name=f"{collection_name}", vectors_config=models.VectorParams( size=1024, distance=models.Distance.COSINE ), ) return Qdrant( client=client, collection_name=collection_name, embeddings=model, ) @staticmethod def _init_online_vector_db( url: str, api_key: str, collection_name: str, model: "HuggingFaceEmbeddings" ): from qdrant_client import QdrantClient """Initialize the Qdrant client that is connected to an online vector store with the given URL and API key. Args: url (str): URL of the Qdrant server. api_key (str): API key for the Qdrant server. """ if api_key is None: if not os.getenv("QDRANT_API_KEY"): raise ValueError("Please provide an api key.") api_key = os.getenv("QDRANT_API_KEY") if url is None: raise ValueError("Please provide a url for the Qdrant server.") try: client = QdrantClient(url=url, api_key=api_key) return QdrantVectorStoreManager._check_create_collection( client=client, collection_name=collection_name, model=model ) except Exception as e: raise ValueError(f"Error occurs when connecting to the server: {e}") @staticmethod def _init_offline_vector_db( vector_store_path: str, collection_name: str, model: "HuggingFaceEmbeddings" ): from qdrant_client import QdrantClient """Initialize the Qdrant client that is connected to an offline vector store with the given vector store folder path. Args: vector_store_path (str): Path to the vector store. """ if vector_store_path is None: raise ValueError("Please provide a folder path.") try: client = QdrantClient(path=vector_store_path) return QdrantVectorStoreManager._check_create_collection( client=client, collection_name=collection_name, model=model ) except Exception as e: raise ValueError(f"Error occurs when loading the vector store: {e}") @staticmethod def create_or_update_vector_store( collection_name: str, vector_db_mode: str, file_path: str, content_column: str, title_column: str = "title", url_column: str = "url", desc_column: str = "description", batch_size: int = 64, chunk_size: int = 500, chunk_overlap: int = 100, vector_store_path: str = None, url: str = None, qdrant_api_key: str = None, embedding_model: str = "BAAI/bge-m3", device: str = "mps", ): from qdrant_client import Document """ Takes a CSV file and adds each row in the CSV file to the Qdrant collection. This function expects each row of the CSV file as a document. The CSV file should have columns for "content", "title", "URL", and "description". Args: collection_name: Name of the Qdrant collection. vector_store_path (str): Path to the directory where the vector store is stored or will be stored. vector_db_mode (str): Mode of the Qdrant vector store (offline or online). file_path (str): Path to the CSV file. content_column (str): Name of the column containing the content. title_column (str): Name of the column containing the title. Default is "title". url_column (str): Name of the column containing the URL. Default is "url". desc_column (str): Name of the column containing the description. Default is "description". batch_size (int): Batch size for adding documents to the collection. chunk_size: Size of each chunk if you need to build the vector store from documents. chunk_overlap: Overlap between chunks if you need to build the vector store from documents. embedding_model: Name of the Hugging Face embedding model. device: Device to run the embeddings model on, can be "mps", "cuda", "cpu". qdrant_api_key: API key for the Qdrant server (Only required if the Qdrant server is online). """ # check if the collection name is provided if collection_name is None: raise ValueError("Please provide a collection name.") model_kwargs = {"device": device} encode_kwargs = {"normalize_embeddings": True} from langchain_huggingface import HuggingFaceEmbeddings model = HuggingFaceEmbeddings( model_name=embedding_model, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs, ) if file_path is None: raise ValueError("Please provide a file path.") # check if the file is a csv file if not file_path.endswith(".csv"): raise ValueError(f"Not valid file format. Please provide a csv file.") if content_column is None: raise ValueError("Please provide the name of the content column.") if url_column is None: raise ValueError("Please provide the name of the url column.") # try to initialize the Qdrant client qdrant = None if vector_db_mode == "online": qdrant = QdrantVectorStoreManager._init_online_vector_db( url=url, api_key=qdrant_api_key, collection_name=collection_name, model=model, ) elif vector_db_mode == "offline": qdrant = QdrantVectorStoreManager._init_offline_vector_db( vector_store_path=vector_store_path, collection_name=collection_name, model=model, ) else: raise ValueError( "Invalid vector_db_mode. Please provide either 'online' or 'offline'." ) if qdrant is None: raise ValueError("Qdrant client is not initialized.") # read the csv file import pandas as pd df = pd.read_csv(file_path) # check that content column exists and url column exists if content_column not in df.columns: raise ValueError( f"Content column {content_column} not found in the csv file." ) if url_column not in df.columns: raise ValueError(f"URL column {url_column} not found in the csv file.") documents = [ Document( page_content=row[content_column], metadata={ "title": row.get(title_column, ""), "url": row[url_column], "description": row.get(desc_column, ""), }, ) for row in df.to_dict(orient="records") ] # split the documents from langchain_text_splitters import RecursiveCharacterTextSplitter text_splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len, add_start_index=True, separators=[ "\n\n", "\n", ".", "\uff0e", # Fullwidth full stop "\u3002", # Ideographic full stop ",", "\uff0c", # Fullwidth comma "\u3001", # Ideographic comma " ", "\u200B", # Zero-width space "", ], ) split_documents = text_splitter.split_documents(documents) # update and save the vector store num_batches = (len(split_documents) + batch_size - 1) // batch_size for i in tqdm(range(num_batches)): start_idx = i * batch_size end_idx = min((i + 1) * batch_size, len(split_documents)) qdrant.add_documents( documents=split_documents[start_idx:end_idx], batch_size=batch_size, ) # close the qdrant client qdrant.client.close() class ArticleTextProcessing: @staticmethod def limit_word_count_preserve_newline(input_string, max_word_count): """ Limit the word count of an input string to a specified maximum, while preserving the integrity of complete lines. The function truncates the input string at the nearest word that does not exceed the maximum word count, ensuring that no partial lines are included in the output. Words are defined as text separated by spaces, and lines are defined as text separated by newline characters. Args: input_string (str): The string to be truncated. This string may contain multiple lines. max_word_count (int): The maximum number of words allowed in the truncated string. Returns: str: The truncated string with word count limited to `max_word_count`, preserving complete lines. """ word_count = 0 limited_string = "" for word in input_string.split("\n"): line_words = word.split() for lw in line_words: if word_count < max_word_count: limited_string += lw + " " word_count += 1 else: break if word_count >= max_word_count: break limited_string = limited_string.strip() + "\n" return limited_string.strip() @staticmethod def remove_citations(s): """ Removes all citations from a given string. Citations are assumed to be in the format of numbers enclosed in square brackets, such as [1], [2], or [1, 2], etc. This function searches for all occurrences of such patterns and removes them, returning the cleaned string. Args: s (str): The string from which citations are to be removed. Returns: str: The string with all citation patterns removed. """ return re.sub(r"\[\d+(?:,\s*\d+)*\]", "", s) @staticmethod def parse_citation_indices(s): """ Extracts citation indexes from the provided content string and returns them as a list of integers. Args: content (str): The content string containing citations in the format [number]. Returns: List[int]: A list of unique citation indexes extracted from the content, in the order they appear. """ matches = re.findall(r"\[\d+\]", s) return [int(index[1:-1]) for index in matches] @staticmethod def remove_uncompleted_sentences_with_citations(text): """ Removes uncompleted sentences and standalone citations from the input text. Sentences are identified by their ending punctuation (.!?), optionally followed by a citation in square brackets (e.g., "[1]"). Grouped citations (e.g., "[1, 2]") are split into individual ones (e.g., "[1] [2]"). Only text up to and including the last complete sentence and its citation is retained. Args: text (str): The input text from which uncompleted sentences and their citations are to be removed. Returns: str: The processed string with uncompleted sentences and standalone citations removed, leaving only complete sentences and their associated citations if present. """ # Convert citations like [1, 2, 3] to [1][2][3]. def replace_with_individual_brackets(match): numbers = match.group(1).split(", ") return " ".join(f"[{n}]" for n in numbers) # Deduplicate and sort individual groups of citations. def deduplicate_group(match): citations = match.group(0) unique_citations = list(set(re.findall(r"\[\d+\]", citations))) sorted_citations = sorted( unique_citations, key=lambda x: int(x.strip("[]")) ) # Return the sorted unique citations as a string return "".join(sorted_citations) text = re.sub(r"\[([0-9, ]+)\]", replace_with_individual_brackets, text) text = re.sub(r"(\[\d+\])+", deduplicate_group, text) # Deprecated: Remove sentence without proper ending punctuation and citations. # Split the text into sentences (including citations). # sentences_with_trailing = re.findall(r'([^.!?]*[.!?].*?)(?=[^.!?]*[.!?]|$)', text) # Filter sentences to ensure they end with a punctuation mark and properly formatted citations # complete_sentences = [] # for sentence in sentences_with_trailing: # # Check if the sentence ends with properly formatted citations # if re.search(r'[.!?]( \[\d+\])*$|^[^.!?]*[.!?]$', sentence.strip()): # complete_sentences.append(sentence.strip()) # combined_sentences = ' '.join(complete_sentences) # Check for and append any complete citations that follow the last sentence # trailing_citations = re.findall(r'(\[\d+\]) ', text[text.rfind(combined_sentences) + len(combined_sentences):]) # if trailing_citations: # combined_sentences += ' '.join(trailing_citations) # Regex pattern to match sentence endings, including optional citation markers. eos_pattern = r"([.!?])\s*(\[\d+\])?\s*" matches = list(re.finditer(eos_pattern, text)) if matches: last_match = matches[-1] text = text[: last_match.end()].strip() return text @staticmethod def clean_up_citation(conv): for turn in conv.dlg_history: if "References:" in turn.agent_utterance: turn.agent_utterance = turn.agent_utterance[ : turn.agent_utterance.find("References:") ] if "Sources:" in turn.agent_utterance: turn.agent_utterance = turn.agent_utterance[ : turn.agent_utterance.find("Sources:") ] turn.agent_utterance = turn.agent_utterance.replace("Answer:", "").strip() try: max_ref_num = max( [int(x) for x in re.findall(r"\[(\d+)\]", turn.agent_utterance)] ) except Exception as e: max_ref_num = 0 if max_ref_num > len(turn.search_results): for i in range(len(turn.search_results), max_ref_num + 1): turn.agent_utterance = turn.agent_utterance.replace(f"[{i}]", "") turn.agent_utterance = ( ArticleTextProcessing.remove_uncompleted_sentences_with_citations( turn.agent_utterance ) ) return conv @staticmethod def clean_up_outline(outline, topic=""): output_lines = [] current_level = 0 # To track the current section level for line in outline.split("\n"): stripped_line = line.strip() if topic != "" and f"# {topic.lower()}" in stripped_line.lower(): output_lines = [] # Check if the line is a section header if stripped_line.startswith("#"): current_level = stripped_line.count("#") output_lines.append(stripped_line) # Check if the line is a bullet point elif stripped_line.startswith("-"): subsection_header = ( "#" * (current_level + 1) + " " + stripped_line[1:].strip() ) output_lines.append(subsection_header) outline = "\n".join(output_lines) # Remove references. outline = re.sub(r"#[#]? See also.*?(?=##|$)", "", outline, flags=re.DOTALL) outline = re.sub(r"#[#]? See Also.*?(?=##|$)", "", outline, flags=re.DOTALL) outline = re.sub(r"#[#]? Notes.*?(?=##|$)", "", outline, flags=re.DOTALL) outline = re.sub(r"#[#]? References.*?(?=##|$)", "", outline, flags=re.DOTALL) outline = re.sub( r"#[#]? External links.*?(?=##|$)", "", outline, flags=re.DOTALL ) outline = re.sub( r"#[#]? External Links.*?(?=##|$)", "", outline, flags=re.DOTALL ) outline = re.sub(r"#[#]? Bibliography.*?(?=##|$)", "", outline, flags=re.DOTALL) outline = re.sub( r"#[#]? Further reading*?(?=##|$)", "", outline, flags=re.DOTALL ) outline = re.sub( r"#[#]? Further Reading*?(?=##|$)", "", outline, flags=re.DOTALL ) outline = re.sub(r"#[#]? Summary.*?(?=##|$)", "", outline, flags=re.DOTALL) outline = re.sub(r"#[#]? Appendices.*?(?=##|$)", "", outline, flags=re.DOTALL) outline = re.sub(r"#[#]? Appendix.*?(?=##|$)", "", outline, flags=re.DOTALL) # clean up citation in outline outline = re.sub(r"\[.*?\]", "", outline) return outline @staticmethod def clean_up_section(text): """Clean up a section: 1. Remove uncompleted sentences (usually due to output token limitation). 2. Deduplicate individual groups of citations. 3. Remove unnecessary summary.""" paragraphs = text.split("\n") output_paragraphs = [] summary_sec_flag = False for p in paragraphs: p = p.strip() if len(p) == 0: continue if not p.startswith("#"): p = ArticleTextProcessing.remove_uncompleted_sentences_with_citations(p) if summary_sec_flag: if p.startswith("#"): summary_sec_flag = False else: continue if ( p.startswith("Overall") or p.startswith("In summary") or p.startswith("In conclusion") ): continue if "# Summary" in p or "# Conclusion" in p: summary_sec_flag = True continue output_paragraphs.append(p) # Join with '\n\n' for markdown format. return "\n\n".join(output_paragraphs) @staticmethod def update_citation_index(s, citation_map): """Update citation index in the string based on the citation map.""" for original_citation in citation_map: s = s.replace( f"[{original_citation}]", f"__PLACEHOLDER_{original_citation}__" ) for original_citation, unify_citation in citation_map.items(): s = s.replace(f"__PLACEHOLDER_{original_citation}__", f"[{unify_citation}]") return s @staticmethod def parse_article_into_dict(input_string): """ Parses a structured text into a nested dictionary. The structure of the text is defined by markdown-like headers (using '#' symbols) to denote sections and subsections. Each section can contain content and further nested subsections. The resulting dictionary captures the hierarchical structure of sections, where each section is represented as a key (the section's title) mapping to a value that is another dictionary. This dictionary contains two keys: - 'content': content of the section - 'subsections': a list of dictionaries, each representing a nested subsection following the same structure. Args: input_string (str): A string containing the structured text to parse. Returns: A dictionary representing contains the section title as the key, and another dictionary as the value, which includes the 'content' and 'subsections' keys as described above. """ lines = input_string.split("\n") lines = [line for line in lines if line.strip()] root = {"content": "", "subsections": {}} current_path = [(root, -1)] # (current_dict, level) for line in lines: if line.startswith("#"): level = line.count("#") title = line.strip("# ").strip() new_section = {"content": "", "subsections": {}} # Pop from stack until find the parent level while current_path and current_path[-1][1] >= level: current_path.pop() # Append new section to the nearest upper level's subsections current_path[-1][0]["subsections"][title] = new_section current_path.append((new_section, level)) else: current_path[-1][0]["content"] += line + "\n" return root["subsections"] class FileIOHelper: @staticmethod def dump_json(obj, file_name, encoding="utf-8"): with open(file_name, "w", encoding=encoding) as fw: json.dump(obj, fw, default=FileIOHelper.handle_non_serializable) @staticmethod def handle_non_serializable(obj): return "non-serializable contents" # mark the non-serializable part @staticmethod def load_json(file_name, encoding="utf-8"): with open(file_name, "r", encoding=encoding) as fr: return json.load(fr) @staticmethod def write_str(s, path): with open(path, "w") as f: f.write(s) @staticmethod def load_str(path): with open(path, "r") as f: return "\n".join(f.readlines()) @staticmethod def dump_pickle(obj, path): with open(path, "wb") as f: pickle.dump(obj, f) @staticmethod def load_pickle(path): with open(path, "rb") as f: return pickle.load(f) class WebPageHelper: """Helper class to process web pages. Acknowledgement: Part of the code is adapted from https://github.com/stanford-oval/WikiChat project. """ def __init__( self, min_char_count: int = 150, snippet_chunk_size: int = 1000, max_thread_num: int = 10, ): """ Args: min_char_count: Minimum character count for the article to be considered valid. snippet_chunk_size: Maximum character count for each snippet. max_thread_num: Maximum number of threads to use for concurrent requests (e.g., downloading webpages). """ self.httpx_client = httpx.Client(verify=False) self.min_char_count = min_char_count self.max_thread_num = max_thread_num self.text_splitter = RecursiveCharacterTextSplitter( chunk_size=snippet_chunk_size, chunk_overlap=0, length_function=len, is_separator_regex=False, separators=[ "\n\n", "\n", ".", "\uff0e", # Fullwidth full stop "\u3002", # Ideographic full stop ",", "\uff0c", # Fullwidth comma "\u3001", # Ideographic comma " ", "\u200B", # Zero-width space "", ], ) def download_webpage(self, url: str): try: res = self.httpx_client.get(url, timeout=4) if res.status_code >= 400: res.raise_for_status() return res.content except httpx.HTTPError as exc: print(f"Error while requesting {exc.request.url!r} - {exc!r}") return None def urls_to_articles(self, urls: List[str]) -> Dict: with concurrent.futures.ThreadPoolExecutor( max_workers=self.max_thread_num ) as executor: htmls = list(executor.map(self.download_webpage, urls)) articles = {} for h, u in zip(htmls, urls): if h is None: continue article_text = extract( h, include_tables=False, include_comments=False, output_format="txt", ) if article_text is not None and len(article_text) > self.min_char_count: articles[u] = {"text": article_text} return articles def urls_to_snippets(self, urls: List[str]) -> Dict: articles = self.urls_to_articles(urls) for u in articles: articles[u]["snippets"] = self.text_splitter.split_text(articles[u]["text"]) return articles def user_input_appropriateness_check(user_input): my_openai_model = LitellmModel( model="azure/gpt-4o-mini", max_tokens=10, temperature=0.0, top_p=0.9, ) if len(user_input.split()) > 20: return "The input is too long. Please make your input topic more concise!" if not re.match(r'^[a-zA-Z0-9\s\-"\,\.?\']*$', user_input): return "The input contains invalid characters. The input should only contain a-z, A-Z, 0-9, space, -/\"/,./?/'." prompt = f"""Here is a topic input into a knowledge curation engine that can write a Wikipedia-like article for the topic. Please judge whether it is appropriate or not for the engine to curate information for this topic based on English search engine. The following types of inputs are inappropriate: 1. Inputs that may be related to illegal, harmful, violent, racist, or sexual purposes. 2. Inputs that are given using languages other than English. Currently, the engine can only support English. 3. Inputs that are related to personal experience or personal information. Currently, the engine can only use information from the search engine. 4. Inputs that are not aimed at topic research or inquiry. For example, asks requiring detailed execution, such as calculations, programming, or specific service searches fall outside the engine's scope of capabilities. If the topic is appropriate for the engine to process, output "Yes."; otherwise, output "No. The input violates reason [1/2/3/4]". User input: {user_input}""" reject_reason_info = { 1: "Sorry, this input may be related to sensitive topics. Please try another topic. " "(Our input filtering uses OpenAI GPT-4o-mini, which may result in false positives. " "We apologize for any inconvenience.)", 2: "Sorry, the current engine can only support English. Please try another topic. " "(Our input filtering uses OpenAI GPT-4o-mini, which may result in false positives. " "We apologize for any inconvenience.)", 3: "Sorry, the current engine cannot process topics related to personal experience. Please try another topic. " "(Our input filtering uses OpenAI GPT-4o-mini, which may result in false positives. " "We apologize for any inconvenience.)", 4: "Sorry, STORM cannot follow arbitrary instruction. Please input a topic you want to learn about. " "(Our input filtering uses OpenAI GPT-4o-mini, which may result in false positives. " "We apologize for any inconvenience.)", } try: response = my_openai_model(prompt)[0].replace("[", "").replace("]", "") if response.startswith("No"): match = regex.search(r"reason\s(\d+)", response) if match: reject_reason = int(match.group(1)) if reject_reason in reject_reason_info: return reject_reason_info[reject_reason] else: return ( "Sorry, the input is inappropriate. Please try another topic!" ) return "Sorry, the input is inappropriate. Please try another topic!" except Exception as e: return "Sorry, the input is inappropriate. Please try another topic!" return "Approved" def purpose_appropriateness_check(user_input): my_openai_model = LitellmModel( model="azure/gpt-4o-mini", max_tokens=10, temperature=0.0, top_p=0.9, ) prompt = f""" Here is a purpose input into a report generation engine that can create a long-form report on any topic of interest. Please judge whether the provided purpose is valid for using this service. Try to judge if given purpose is non-sense like random words or just try to get around the sanity check. You should not make the rule too strict. If the purpose is valid, output "Yes."; otherwise, output "No" followed by reason. User input: {user_input} """ try: response = my_openai_model(prompt)[0].replace("[", "").replace("]", "") if response.startswith("No"): return "Please provide a more detailed explanation on your purpose of requesting this article." except Exception as e: return "Please provide a more detailed explanation on your purpose of requesting this article." return "Approved" ================================================ FILE: requirements.txt ================================================ dspy_ai==2.4.9 wikipedia==1.4.0 sentence-transformers toml langchain-text-splitters trafilatura langchain-huggingface qdrant-client langchain-qdrant numpy litellm diskcache ================================================ FILE: setup.py ================================================ import re from setuptools import setup, find_packages # Read the content of the README file with open("README.md", encoding="utf-8") as f: long_description = f.read() # Remove p tags. pattern = re.compile(r".*?

    ", re.DOTALL) long_description = re.sub(pattern, "", long_description) # Read the content of the requirements.txt file with open("requirements.txt", encoding="utf-8") as f: requirements = f.read().splitlines() setup( name="knowledge-storm", version="1.1.1", author="Yijia Shao, Yucheng Jiang", author_email="shaoyj@stanford.edu, yuchengj@stanford.edu", description="STORM: A language model-powered knowledge curation engine.", long_description=long_description, long_description_content_type="text/markdown", url="https://github.com/stanford-oval/storm", license="MIT License", packages=find_packages(), classifiers=[ "Development Status :: 3 - Alpha", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", ], python_requires=">=3.10", install_requires=requirements, )