app.py

import os
import pandas as pd
import streamlit as st
from datetime import datetime
from groq import Groq
from logic import LLMClient, CodeProcessor
from batch_code_logic_csv import csv_read_batch_code
import zipfile
import io
import markdown2
import pdfkit
import tempfile

client = Groq(api_key=os.getenv("GROQ_API_KEY"))
llm_obj = LLMClient(client)
processor = CodeProcessor(llm_obj)

st.title("Code Analysis with LLMs")

st.sidebar.title("Input Options")

code_input_method = st.sidebar.radio("How would you like to provide your code?",
                                     ("Upload CSV file", "Upload Code File"))

code_dict = {}

if code_input_method == "Upload CSV file":
    uploaded_file = st.sidebar.file_uploader("Upload your CSV/Excel file", type=["csv", "xlsx"])

    if uploaded_file is not None:
        dataframe = pd.read_csv(uploaded_file)
        code_dict = csv_read_batch_code(dataframe)

elif code_input_method == "Upload Code File":
    uploaded_file = st.sidebar.file_uploader("Upload your code file", type=["py", "txt"])
    if uploaded_file is not None:
        code_text = uploaded_file.read().decode("utf-8")
        code_dict = {"single_code": code_text}

model_choice = st.sidebar.selectbox("Select LLM Model",
                                    ["llama-3.2-90b-text-preview", "llama-3.2-90b-text-preview", "llama3-8b-8192"])

if code_dict:
    unique_key = st.sidebar.selectbox("Select a Key for Analysis", list(code_dict.keys()))

    if st.sidebar.button("Analyze Code") and unique_key:
        code_text = code_dict[unique_key]
        markdown_output = processor.process_code(code_text, model_choice)

        with st.expander(f"Analysis for {unique_key}"):
            st.markdown(markdown_output)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        with tempfile.NamedTemporaryFile(delete=False, suffix=".md") as tmp_md_file:
            tmp_md_file.write(markdown_output.encode('utf-8'))
            tmp_md_file_path = tmp_md_file.name

        with open(tmp_md_file_path, "r") as md_file:
            st.download_button(
                label=f"Download {unique_key} Result as Markdown",
                data=md_file,
                file_name=f"code_analysis_{unique_key}_{timestamp}.md",
                mime="text/markdown"
            )

        html_output = markdown2.markdown(markdown_output)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf_file:
            pdf_file_path = tmp_pdf_file.name
            pdfkit.from_string(html_output, pdf_file_path)

        with open(pdf_file_path, "rb") as pdf_file:
            st.download_button(
                label=f"Download {unique_key} Result as PDF",
                data=pdf_file,
                file_name=f"code_analysis_{unique_key}_{timestamp}.pdf",
                mime="application/pdf"
            )

    if st.sidebar.button("Batch Predict"):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        all_markdowns = {}

        for key, code_text in code_dict.items():
            markdown_output = processor.process_code(code_text, model_choice)
            all_markdowns[key] = markdown_output

            with st.expander(f"Analysis for {key}"):
                st.markdown(markdown_output)

            with tempfile.NamedTemporaryFile(delete=False, suffix=".md") as tmp_md_file:
                tmp_md_file.write(markdown_output.encode('utf-8'))
                tmp_md_file_path = tmp_md_file.name

            with open(tmp_md_file_path, "r") as md_file:
                st.download_button(
                    label=f"Download {key} Result as Markdown",
                    data=md_file,
                    file_name=f"code_analysis_{key}_{timestamp}.md",
                    mime="text/markdown"
                )

            html_output = markdown2.markdown(markdown_output)
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf_file:
                pdf_file_path = tmp_pdf_file.name
                pdfkit.from_string(html_output, pdf_file_path)

            with open(pdf_file_path, "rb") as pdf_file:
                st.download_button(
                    label=f"Download {key} Result as PDF",
                    data=pdf_file,
                    file_name=f"code_analysis_{key}_{timestamp}.pdf",
                    mime="application/pdf"
                )

        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
            for key, markdown_output in all_markdowns.items():
                zip_file.writestr(f"code_analysis_{key}_{timestamp}.md", markdown_output)

        st.download_button(
            label="Download All as Zip",
            data=zip_buffer.getvalue(),
            file_name=f"code_analysis_batch_{timestamp}.zip",
            mime="application/zip"
        )
else:
    st.write("Please upload your file to analyze.")