brunogreen25/AICodingAnalyst

1import streamlit as st
2from prompt_generation import CustomLLM
3from utils import PAGE
4from io import StringIO
5from dotenv import load_dotenv
6load_dotenv()
7
8
9def init():
10    if 'current_page' not in st.session_state:
11        st.session_state.current_page = PAGE.MAIN
12        st.session_state.llm = CustomLLM(temperature=0.)  # Other models: ["gpt-3.5-turbo-0125", "gpt-4", "gpt-4o"]
13
14    # Setting page title and header
15    st.set_page_config(page_title="AICodingAnalyst", page_icon="💻", layout="centered", initial_sidebar_state="collapsed")
16    st.markdown("<h1 style='text-align: center;'>AI Coding Analyst</h1>", unsafe_allow_html=True)
17
18
19def page_main():
20    # Header
21    st.header("Choose what to do with your code:")
22    col1, col2 = st.columns(2)
23
24    with col1:
25        if st.button("Translate to a new programming language"):
26            st.session_state.current_page = PAGE.TRANSLATE
27            st.rerun()
28    with col2:
29        if st.button("Review and optimize my code"):
30            st.session_state.current_page = PAGE.REVIEW
31            st.rerun()
32
33
34def page_translate():
35    # Header
36    if st.button(":back: Main Page"):
37        st.session_state.current_page = PAGE.MAIN
38        st.rerun()
39    st.header("Choose the source code and language")
40
41    # Upload code file
42    uploaded_file = st.file_uploader("Upload your code", accept_multiple_files=False)
43
44    # Select a language to translate to
45    language = st.text_input("Language", value="Powershell")
46
47    # button
48    button_translate = st.button("Translate")
49
50    if button_translate:
51        # Read code
52        stringio = StringIO(uploaded_file.getvalue().decode('utf-8'))
53        code = stringio.read()
54
55        # Get the output
56        with st.spinner("Translating the code..."):
57            translated_code = st.session_state.llm.translate_code(code, language)
58            st.write(translated_code)
59
60
61def page_review():
62    # Header
63    if st.button(":back: Main Page"):
64        st.session_state.current_page = PAGE.MAIN
65        st.rerun()
66    st.header("Code Reviewer")
67
68    # Upload code file
69    uploaded_file = st.file_uploader("Upload your code", accept_multiple_files=False)
70
71    # Button
72    button_review = st.button("Review")
73
74    if button_review:
75        # Read code
76        stringio = StringIO(uploaded_file.getvalue().decode('utf-8'))
77        code = stringio.read()
78
79        # Ge the output
80        with st.spinner("Reviewing the code..."):
81            new_code, comments = st.session_state.llm.review_code(code)
82            st.code(new_code)
83            st.write(comments)
84
85
86# Main structure
87init()
88
89# Page selector
90match st.session_state.current_page:
91    case PAGE.MAIN:
92        page_main()
93    case PAGE.TRANSLATE:
94        page_translate()
95    case PAGE.REVIEW:
96        page_review()
97