├── requirements.txt ├── doc ├── revise.png ├── whole_revise.png └── initial_prompt.png ├── CODE_OF_CONDUCT.md ├── config.py ├── prompts.py ├── LICENSE ├── README.md ├── llm.py ├── CONTRIBUTING.md └── main.py /requirements.txt: -------------------------------------------------------------------------------- 1 | boto3==1.28.70 2 | streamlit==1.30.0 3 | -------------------------------------------------------------------------------- /doc/revise.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/ai-writer/main/doc/revise.png -------------------------------------------------------------------------------- /doc/whole_revise.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/ai-writer/main/doc/whole_revise.png -------------------------------------------------------------------------------- /doc/initial_prompt.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-samples/ai-writer/main/doc/initial_prompt.png -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | ## Code of Conduct 2 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 3 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 4 | opensource-codeofconduct@amazon.com with any additional questions or comments. 5 | -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | AWS_REGION="us-west-2" 2 | MODEL_ID="anthropic.claude-3-sonnet-20240229-v1:0" 3 | MODEL_PARAMETERS={ 4 | "anthropic_version": "bedrock-2023-05-31", 5 | "max_tokens": 4096, 6 | "temperature": 1.0, 7 | } 8 | 9 | ## Example for Claude 2 10 | # AWS_REGION="eu-central-1" 11 | # MODEL_ID="anthropic.claude-v2" 12 | # MODEL_PARAMETERS={ 13 | # "max_tokens_to_sample": 2048, 14 | # "temperature": 1.0, 15 | # "top_k": 250, 16 | # "top_p": 1, 17 | # "stop_sequences": [ "\\n\\nHuman:" ], 18 | # "anthropic_version": "bedrock-2023-05-31" 19 | # } 20 | -------------------------------------------------------------------------------- /prompts.py: -------------------------------------------------------------------------------- 1 | # TODO: encapsulate this into the llm 2 | 3 | def writing_prompt(user_prompt): 4 | return f"""Write {user_prompt}. Only show me what you write, do NOT say something like "Here is an article:" or "Here is a story" in the beginning.""" 5 | 6 | def revise_prompt(user_prompt, current_paragraph): 7 | return f""" 8 | Revise the following paragraph this way: {user_prompt}. Only output the revised paragraph. Do NOT say something like "Here is the revised paragraph:" in the beginning. 9 | --- 10 | {current_paragraph} 11 | --- 12 | """ 13 | 14 | def overall_revise_prompt(user_prompt, current_paragraph): 15 | return f""" 16 | Revise the whole article this way: {user_prompt}. Output the whole article, including the paragraphs that have not changed. Do NOT say something like "Here is the revised article:" in the beginnging. 17 | --- 18 | {current_paragraph} 19 | --- 20 | """ 21 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT No Attribution 2 | 3 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | this software and associated documentation files (the "Software"), to deal in 7 | the Software without restriction, including without limitation the rights to 8 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | the Software, and to permit persons to whom the Software is furnished to do so. 10 | 11 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 12 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 13 | FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 14 | COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 15 | IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 16 | CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 17 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AI Writer 2 | 3 | This is a Streamlit application that allows users to generate and iteratively revise AI-written articles. 4 | 5 | ## Functionality 6 | 7 | The application has the following functionality: 8 | 9 | - Users can provide a prompt to have the AI generate an initial article draft. The generated text is split into paragraphs. ![prompt screenshot](doc/initial_prompt.png) 10 | 11 | - The drafted article paragraphs are displayed in a list. Users can click on each paragraph to edit it in a textarea. 12 | 13 | - Below each paragraph textarea is a "Revise" button. Users can provide revise instructions and click this to have the AI revise just that paragraph. ![revise screenshot](doc/revise.png) 14 | 15 | - There is an option to provide revise instructions for the entire article, which will invoke the AI to revise all paragraphs. ![whole revise screenshot](doc/whole_revise.png) 16 | 17 | 18 | - The edited article can be copied out of a final textarea field. 19 | 20 | 21 | ## Setup 22 | 23 | Requirements: 24 | 25 | - Python 3 26 | - An AWS account with [Amazon Bedrock access](https://docs.aws.amazon.com/bedrock/latest/userguide/security_iam_id-based-policy-examples.html#security_iam_id-based-policy-examples-console) 27 | - Make sure you have the required [IAM permission](https://docs.aws.amazon.com/bedrock/latest/userguide/security_iam_id-based-policy-examples.html#security_iam_id-based-policy-examples-console) and accepted the model [EULAs](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html) 28 | - Follow [this guide](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html) to set up the AWS SDK credentails. 29 | 30 | ``` 31 | pip install -r requirements.txt 32 | ``` 33 | 34 | Start the application: 35 | 36 | ``` 37 | streamlit run main.py --server.port 8080 38 | ``` 39 | 40 | Open a browser and navigate to http://localhost:8080 to get started. 41 | 42 | 43 | ## Contributing 44 | 45 | See [CONTRIBUTING](CONTRIBUTING.md) for more infromation 46 | 47 | ## License 48 | 49 | See [LICENSE](LICENSE) 50 | -------------------------------------------------------------------------------- /llm.py: -------------------------------------------------------------------------------- 1 | import boto3 2 | import json 3 | import config 4 | 5 | def invoke_llm(prompt): 6 | if config.MODEL_ID == 'anthropic.claude-3-sonnet-20240229-v1:0': 7 | return invoke_claude3sonnet(prompt) 8 | elif config.MODEL_ID == 'anthropic.claude-v2': 9 | return invoke_claude2(prompt) 10 | else: 11 | raise Exception("Invalid model") 12 | return None 13 | 14 | def invoke_claude3sonnet(prompt): 15 | client = boto3.client('bedrock-runtime', region_name=config.AWS_REGION) 16 | 17 | try: 18 | response = client.invoke_model( 19 | body = json.dumps({ 20 | "messages": [ 21 | { 22 | "role": "user", 23 | "content": [ 24 | { 25 | "type": "text", 26 | "text": prompt 27 | } 28 | ] 29 | } 30 | ], 31 | **config.MODEL_PARAMETERS 32 | }), 33 | modelId = config.MODEL_ID, 34 | contentType = "application/json", 35 | accept = "application/json" 36 | ) 37 | response_text = response['body'].read().decode('utf8').strip() 38 | response_json = json.loads(response_text) 39 | print(response_json) 40 | completion = response_json['content'][0]['text'] 41 | return completion 42 | except Exception as e: 43 | print('Error invoking endpoint') 44 | print(e) 45 | raise Exception("Error invoking LLM. See Python CLI output for the full error message.") 46 | 47 | def invoke_claude2(prompt): 48 | client = boto3.client('bedrock-runtime', region_name=config.AWS_REGION) 49 | 50 | claude2_prompt = f""" 51 | Human: {prompt} 52 | 53 | Assistant:""" 54 | 55 | try: 56 | response = client.invoke_model( 57 | body = bytes(json.dumps({ 58 | "prompt": claude2_prompt, 59 | **config.MODEL_PARAMETERS 60 | }), 'utf-8'), 61 | modelId = config.MODEL_ID, 62 | contentType = "application/json", 63 | accept = "application/json" 64 | ) 65 | response_text = response['body'].read().decode('utf8').strip() 66 | response_json = json.loads(response_text) 67 | completion = response_json['completion'] 68 | return completion 69 | except Exception as e: 70 | print('Error invoking endpoint') 71 | print(e) 72 | raise Exception("Error invoking LLM") -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing Guidelines 2 | 3 | Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional 4 | documentation, we greatly value feedback and contributions from our community. 5 | 6 | Please read through this document before submitting any issues or pull requests to ensure we have all the necessary 7 | information to effectively respond to your bug report or contribution. 8 | 9 | 10 | ## Reporting Bugs/Feature Requests 11 | 12 | We welcome you to use the GitHub issue tracker to report bugs or suggest features. 13 | 14 | When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already 15 | reported the issue. Please try to include as much information as you can. Details like these are incredibly useful: 16 | 17 | * A reproducible test case or series of steps 18 | * The version of our code being used 19 | * Any modifications you've made relevant to the bug 20 | * Anything unusual about your environment or deployment 21 | 22 | 23 | ## Contributing via Pull Requests 24 | Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that: 25 | 26 | 1. You are working against the latest source on the *main* branch. 27 | 2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already. 28 | 3. You open an issue to discuss any significant work - we would hate for your time to be wasted. 29 | 30 | To send us a pull request, please: 31 | 32 | 1. Fork the repository. 33 | 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change. 34 | 3. Ensure local tests pass. 35 | 4. Commit to your fork using clear commit messages. 36 | 5. Send us a pull request, answering any default questions in the pull request interface. 37 | 6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation. 38 | 39 | GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and 40 | [creating a pull request](https://help.github.com/articles/creating-a-pull-request/). 41 | 42 | 43 | ## Finding contributions to work on 44 | Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start. 45 | 46 | 47 | ## Code of Conduct 48 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 49 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 50 | opensource-codeofconduct@amazon.com with any additional questions or comments. 51 | 52 | 53 | ## Security issue notifications 54 | If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue. 55 | 56 | 57 | ## Licensing 58 | 59 | See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution. 60 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | # a streamlit application that shows a list of text paragraphs 2 | # when click on each text paragraph, it becomes a text area to allow editing, also a input box will show below the text paragraph 3 | 4 | 5 | import streamlit as st 6 | import sys 7 | import boto3 8 | import json 9 | import re 10 | 11 | from prompts import * 12 | from llm import invoke_llm 13 | import config 14 | 15 | def split_paragraphs(full_article): 16 | # split the full_article into an array of paragraphs by a new line, but keep markdown code blocks wrapped by three backquotes in one paragraph 17 | paragraphs = full_article.splitlines(True) 18 | output = [] 19 | in_code_block = False 20 | for paragraph in paragraphs: 21 | if in_code_block: 22 | if "```" in paragraph: 23 | in_code_block = False 24 | output[-1] += paragraph 25 | continue 26 | elif paragraph.startswith("```"): 27 | in_code_block = True 28 | output.append(paragraph) 29 | continue 30 | elif paragraph.strip() != "": 31 | output.append(paragraph) 32 | return output 33 | 34 | if "article" not in st.session_state: 35 | st.session_state["article"] = [] 36 | 37 | if "editing_idx" not in st.session_state: 38 | st.session_state["editing_idx"] = None 39 | 40 | def main(): 41 | st.title("AI Writer") 42 | 43 | is_starting_over = len(st.session_state["article"]) == 0 44 | 45 | with st.expander("Start over", expanded = is_starting_over): 46 | article_prompt = st.text_area("What would you like to write?", placeholder="A short story about an unicorn, a happy birthday email") 47 | 48 | if st.button("Write"): 49 | # Using this instead of using button.disabled because when you enter something into the text_area, 50 | # you need to unfocus the text_area before the button's state updates. This is bad UX. 51 | if article_prompt.strip() == "": 52 | st.error("Please enter a prompt") 53 | return 54 | raw_article = invoke_llm(writing_prompt(article_prompt)) 55 | st.session_state["article"] = split_paragraphs(raw_article) 56 | st.rerun() 57 | 58 | with st.expander("Revise the whole article", expanded = not is_starting_over): 59 | overall_revise_instruction = st.text_area("How would you like to revise the whole article?", placeholder="Change from third-person to first-person. Make it longer.") 60 | if st.button("Revise"): 61 | if st.session_state["article"] == []: 62 | st.error("Please write something first") 63 | return 64 | if overall_revise_instruction.strip() == "": 65 | st.error("Please enter a prompt") 66 | return 67 | revised_article = invoke_llm(overall_revise_prompt(overall_revise_instruction, "\n\n".join(st.session_state["article"]))) 68 | st.session_state["article"] = split_paragraphs(revised_article) 69 | 70 | with st.expander("Import", expanded = False): 71 | import_article = st.text_area("Paste an existing article here", placeholder="Paste an existing article here") 72 | if st.button("Import"): 73 | if import_article.strip() == "": 74 | st.error("Please paste a valid article") 75 | return 76 | st.session_state["article"] = split_paragraphs(import_article) 77 | st.session_state["editing_idx"] = None 78 | st.session_state["article"] = split_paragraphs("\n".join(st.session_state["article"])) # One paragraph might be edited to two 79 | st.rerun() 80 | 81 | # Add a devider 82 | st.markdown("---") 83 | 84 | for idx, paragraph in enumerate(st.session_state["article"]): 85 | if idx == st.session_state["editing_idx"]: 86 | st.markdown("---") 87 | text_col, action_col = st.columns([0.8, 0.2]) 88 | with text_col: 89 | editing_text_area = st.text_area("(Editing)", 90 | value=paragraph, 91 | #on_change=update_paragraph 92 | ) 93 | with action_col: 94 | if editing_text_area != st.session_state["article"][idx] or st.button("Save", key=f"edit-{idx}"): 95 | st.session_state["article"][idx] = editing_text_area 96 | st.session_state["editing_idx"] = None 97 | st.session_state["article"] = split_paragraphs("\n".join(st.session_state["article"])) # One paragraph might be edited to two 98 | st.rerun() 99 | revise_instruction= st.text_area("How would you like to revise this paragraph?", placeholder="Make the tone softer") 100 | if st.button("Revise", key=f"revise-{idx}"): 101 | if revise_instruction.strip() == "": 102 | st.error("Please enter a prompt") 103 | return 104 | revised_paragraph = invoke_llm(revise_prompt(revise_instruction, editing_text_area)) 105 | st.session_state["article"][idx] = revised_paragraph 106 | st.session_state["article"] = split_paragraphs("\n".join(st.session_state["article"])) # One paragraph might be edited to two 107 | st.rerun() 108 | st.markdown("---") 109 | 110 | 111 | else: 112 | text_col, action_col_1, action_col_2 = st.columns([0.8, 0.1, 0.12]) 113 | with text_col: 114 | paragraph 115 | # when the paragraph is clicked, turn the text into a text area with the paragraph as placeholder 116 | with action_col_1: 117 | if st.button("Edit", key=f"edit-{idx}"): 118 | st.session_state["editing_idx"] = idx 119 | st.rerun() 120 | with action_col_2: 121 | if st.button("Delete", key=f"delete-{idx}"): 122 | st.session_state["article"].pop(idx) 123 | st.rerun() 124 | 125 | st.markdown('---') 126 | st.markdown("## Copy the article below") 127 | st.text_area("Finished article", value="\n".join(st.session_state["article"])) 128 | 129 | if __name__ == "__main__": 130 | main() --------------------------------------------------------------------------------