CintraAI/code-chunker

1import streamlit as st2import json3import os4from Chunker import CodeChunker5 6# Set Streamlit page config at the very beginning7st.set_page_config(page_title="Cintra Code Chunker", layout="wide")8 9# Function to load JSON data10def load_json_file(file_path):11    with open(file_path, 'r') as file:12        return json.load(file)13 14# Function to read code from an uploaded file15def read_code_from_file(uploaded_file):16    return uploaded_file.getvalue().decode("utf-8")17 18st.link_button('Contribute on GitHub', 'https://github.com/CintraAI/code-chunker', help=None, type="secondary", disabled=False, use_container_width=False)19 20json_file_path = os.path.join(os.path.dirname(__file__), 'mock_codefiles.json')21code_files_data = load_json_file(json_file_path)22 23# Extract filenames and contents24code_files = list(code_files_data.keys())25 26st.title('Cintra Code Chunker')27 28selection_col, upload_col = st.columns(2)29with selection_col:30    # File selection dropdown31    selected_file_name = st.selectbox("Select an example code file", code_files)32 33with upload_col:34    # File upload35    uploaded_file = st.file_uploader("Or upload your code file", type=['py', 'js', 'css', 'jsx'])36 37# Determine the content and file extension based on selection or upload38if uploaded_file is not None:39    code_content = read_code_from_file(uploaded_file)40    file_extension = uploaded_file.name.split('.')[-1]41else:42    code_content = code_files_data.get(selected_file_name, "")43    file_extension = selected_file_name.split('.')[-1] if selected_file_name else None44 45# Determine the language for syntax highlighting46def get_language_by_extension(file_extension):47    if file_extension in ['py', 'python']:48        return 'python'49    elif file_extension in ['js', 'jsx', 'javascript']:50        return 'javascript'51    elif file_extension == 'css':52        return 'css'53    elif file_extension in ['ts', 'typescript', 'tsx']:54        return 'typescript'55    elif file_extension in ['rb', 'ruby']:56        return 'ruby'57    elif file_extension == 'php':58        return 'php'59    elif file_extension == 'go':60        return 'go'61    else:62        return None63 64language = get_language_by_extension(file_extension)65 66st.write("""67### Choose Chunk Size Target""")68token_chunk_size = st.number_input('Target Chunk Size Target', min_value=5, max_value=1000, value=25, help="The token limit guides the chunk size in tokens (tiktoken, gpt-4), aiming for readability without enforcing a strict upper limit.")69 70with st.expander("Learn more about the chunk size target"):71    st.markdown("""72The `token_limit` parameter in the `chunk` function serves as a guideline to optimize the size of code chunks produced. It is not a hard limit but rather an ideal target, attempting to achieve a balance between chunk size and maintaining logical coherence within the code.73 74- **Adherence to Logical Breakpoints:** The chunking logic respects logical breakpoints in the code, ensuring that chunks are coherent and maintain readability.75- **Flexibility in Chunk Size:** Chunks might be slightly smaller or larger than the specified `token_limit` to avoid breaking the code in the middle of logical sections.76- **Handling Final Chunks:** The last chunk of code captures any remaining code, which may vary significantly in size depending on the remaining code's structure.77 78This approach allows for flexibility in how code is segmented into chunks, emphasizing the balance between readable, logical code segments and size constraints.79    """)80 81original_col, chunked_col = st.columns(2)82 83with original_col:84    st.subheader('Original File')85    st.code(code_content, language=language)86 87# Initialize the code chunker88code_chunker = CodeChunker(file_extension=file_extension)89 90# Chunk the code content91chunked_code_dict = code_chunker.chunk(code_content, token_chunk_size)92 93with chunked_col:94    st.subheader('Chunked Code')95    for chunk_key, chunk_code in chunked_code_dict.items():96        st.code(chunk_code, language=language)