Upload folder using huggingface_hub

.docker/Dockerfile

@@ -0,0 +1,20 @@
+# Use an official Python image
+FROM python:3.13
+
+# Set the working directory
+WORKDIR /app
+
+# Copy only requirements first (for better caching)
+COPY requirements.txt .
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application code
+COPY . .
+
+# Expose the port Gradio runs on
+EXPOSE 7860
+
+# Command to run the Gradio app
+CMD ["python", "interface.py"]

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+models

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

README.md

@@ -1,12 +1,7 @@
----
-title: Code Repair W LLMs
-emoji: 🏢
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 5.16.2
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Code-Repair-w-LLMs
+app_file: main.py
+sdk: gradio
+sdk_version: 5.6.0
+---
+frontend for code repair w llms

api/initiate_pipeline_call.py

@@ -0,0 +1,34 @@
+import requests
+import os
+import mimetypes
+
+API_URL = "http://localhost:8000/api/initiate_pipeline"
+
+def initiate_pipeline_call(files, pipeline_steps):
+    """
+    Initiate the pipeline with multiple files.
+    
+    Args:
+        files: List of tuples containing file information (name, content, type)
+        pipeline_steps: Integer representing selected pipeline steps
+    """
+
+    file_data = []
+    
+    for file in files:
+        mime_type = mimetypes.guess_type(file.name)[0] or "application/octet-stream"
+        
+        # Ensure we send the actual file content
+        with open(file.name, "rb") as f:
+            file_data.append(("files", (file.name, f.read(), mime_type)))  
+
+    try:
+        response = requests.post(
+            f"{API_URL}?pipeline_steps={pipeline_steps}", 
+            files=file_data
+        )
+        
+        response.raise_for_status()  # Raise exception for bad status codes
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        return f"Error: {str(e)}"

components/bug_finding.py

@@ -0,0 +1,8 @@
+import gradio as gr
+from components.model_selection import create_model_selection_dropdown
+
+def create_bug_finding_tab(choices):
+    with gr.Column():
+        create_model_selection_dropdown(choices)
+        gr.Textbox(label="Output", interactive=False)
+        gr.Code(label="File Contents", language="python", interactive=False, elem_classes=["fixed-height"])

components/chat.py

@@ -0,0 +1,13 @@
+import gradio as gr
+
+def create_chat_controls():
+    with gr.Row():
+        with gr.Column(scale=10):
+            msg = gr.Textbox(label="Prompt", placeholder="Enter prompt")
+
+    with gr.Row():
+        with gr.Column(min_width=0, scale=10):
+            submit_button = gr.Button("Submit Prompt")
+
+    return msg, submit_button
+

components/file_upload.py

@@ -0,0 +1,10 @@
+import gradio as gr
+
+def create_file_upload_section():
+    with gr.Row(equal_height=True):  # Ensures both sections are the same height
+        with gr.Column(scale=1):
+            file_input = gr.File(label="Upload Codebase/Single File", file_types=[".py", ".java", ".c", ".cpp"])
+        with gr.Column(scale=3):
+            file_content = gr.Code(label="File Contents", language="python", interactive=False, elem_classes=["fixed-height"])
+    
+    return file_input, file_content

components/front_page.py

@@ -0,0 +1,66 @@
+import gradio as gr
+import utils.interface_utils as iutils
+
+from api.initiate_pipeline_call import initiate_pipeline_call
+from utils.convert_steps import convert_steps
+from utils.storage import save_jwt_to_session
+from components.interface import create_interface
+from components.model_selection import create_model_selection_dropdown
+
+def handle_initiate_pipeline(files, selected_steps, initial_prompt):
+    """Processes the file upload and initiates the pipeline."""
+    if not files:
+        gr.Warning("Please upload a file/codebase.")
+        return gr.update(visible=True), gr.update(visible=False)
+    
+    binary_steps = convert_steps(selected_steps)
+
+    # TODO: Do something with prompt
+    
+    response = initiate_pipeline_call(files, binary_steps)
+    '''
+        Delete print below in production
+    '''
+    print(response)
+    save_jwt_to_session(response)  
+
+    return gr.update(visible=False), gr.update(visible=True)
+
+def create_full_ui():
+    """Creates the full UI with a front page and the main interface."""
+    with gr.Blocks(theme=iutils.custom_theme(), css="style.css") as app:
+        # TODO: Replace with actual choices
+        choices = ["ChatGPT", "Claude"]
+
+        # Front Page UI
+        with gr.Column(visible=True) as front_page:
+            gr.Markdown("# Welcome to Code Repair with LLMs!")
+            gr.Markdown("## Upload your files/codebase. All steps of the pipeline will be executed on the initial run.")
+
+            file_input = gr.Files(
+                label="Upload Multiple Files / Codebase",
+                file_types=[".py", ".java", ".c", ".cpp"]
+            )
+            gr.Markdown("Python, Java, C++, and .zip files accepted.")
+
+
+            steps = ["Bug Finding", "Pattern Matching", "Patch Generation", "Patch Validation"]
+            checkboxes = gr.CheckboxGroup(steps, label="Select Desired Steps", value=steps, interactive=False)
+
+            create_model_selection_dropdown(choices=choices)
+
+            msg = gr.Textbox(label="Prompt", placeholder="Enter prompt (Optional)")
+            initiate_button = gr.Button("Initiate Pipeline")
+
+        # Main Interface UI (Initially Hidden)
+        with gr.Row(visible=False) as main_interface:
+            interface = create_interface()
+
+        # Button click should happen inside the context
+        initiate_button.click(
+            fn=handle_initiate_pipeline,
+            inputs=[file_input, checkboxes, msg],
+            outputs=[front_page, main_interface]
+        )
+
+    return app

components/interface.py

@@ -0,0 +1,51 @@
+import gradio as gr
+import os
+import utils.interface_utils as iutils
+from components.chat import create_chat_controls
+from components.bug_finding import create_bug_finding_tab
+from components.pattern_matching import create_pattern_matching_tab
+from components.patch_generation import create_patch_generation_tab
+from components.patch_validation import create_patch_validation_tab
+from components.file_upload import create_file_upload_section
+from components.model_selection import create_model_selection_dropdown
+
+
+def create_interface():
+    # TODO: Replace with actual choices
+    choices = ["ChatGPT", "Claude"]
+
+    with gr.Blocks(theme=iutils.custom_theme(), css=iutils.custom_css()) as interface:
+        with gr.Column():
+            # Header with Logo
+            with gr.Row():
+                gr.Markdown("# Code Repair with LLMs")  
+                gr.Markdown("")  # Filler to push logo to the right
+                gr.Image("logo.png", label="Logo", show_download_button=False, show_fullscreen_button=False, show_label=False, height=250, width=450, container=False)
+
+            # Tabs for Different Features
+            with gr.Tab("Chat"):
+                with gr.Column():
+                    steps = ["Bug Finding", "Pattern Matching", "Patch Generation", "Patch Validation"]
+                    checkboxes = gr.CheckboxGroup(steps, label="Select Desired Steps", value=steps, interactive=True)
+                    create_model_selection_dropdown(choices)
+                    with gr.Column():
+                        chatbot = gr.Chatbot(value=None, type="messages", show_label=True, show_share_button=False)
+                        msg, submit_button = create_chat_controls()
+
+
+            with gr.Tab("Bug Finding"):
+                create_bug_finding_tab(choices)
+
+            with gr.Tab("Pattern Matching"):
+                create_pattern_matching_tab(choices)
+
+            with gr.Tab("Patch Generation"):
+                create_patch_generation_tab(choices)
+
+            with gr.Tab("Patch Validation"):
+                create_patch_validation_tab(choices)
+
+        # TODO: need new function to handle future pipeline calls
+        # msg.submit(fn=handle_initiate_pipeline, inputs=[files, checkboxes, msg], outputs=[chatbot]) # clone for submit_button
+
+    return interface

components/model_selection.py

@@ -0,0 +1,4 @@
+import gradio as gr
+
+def create_model_selection_dropdown(choices):
+    return gr.Dropdown(label="Model Selection", choices=choices, interactive=True)

components/patch_generation.py

@@ -0,0 +1,8 @@
+import gradio as gr
+from components.model_selection import create_model_selection_dropdown
+
+def create_patch_generation_tab(choices):
+    with gr.Column():
+        create_model_selection_dropdown(choices)
+        gr.Textbox(label="Output", interactive=False)
+        gr.Code(label="File Contents", language="python", interactive=False, elem_classes=["fixed-height"])

components/patch_validation.py

@@ -0,0 +1,8 @@
+import gradio as gr
+from components.model_selection import create_model_selection_dropdown
+
+def create_patch_validation_tab(choices):
+    with gr.Column():
+        create_model_selection_dropdown(choices)
+        gr.Textbox(label="Output", interactive=False)
+        gr.Code(label="File Contents", language="python", interactive=False, elem_classes=["fixed-height"])

components/pattern_matching.py

@@ -0,0 +1,8 @@
+import gradio as gr
+from components.model_selection import create_model_selection_dropdown
+
+def create_pattern_matching_tab(choices):
+    with gr.Column():
+        create_model_selection_dropdown(choices)
+        gr.Textbox(label="Output", interactive=False)
+        gr.Code(label="File Contents", language="python", interactive=False, elem_classes=["fixed-height"])

docker-compose.yml

@@ -0,0 +1,14 @@
+version: "3.8"
+
+services:
+  gradio_app:
+    build:
+      context: .  # Use the root as context
+      dockerfile: .docker/Dockerfile  # Specify the Dockerfile inside .docker
+    ports:
+      - "7860:7860"
+    volumes:
+      - .:/app
+    environment:
+      - PYTHONUNBUFFERED=1
+    restart: unless-stopped

main.py

@@ -0,0 +1,10 @@
+import gradio as gr
+from components.front_page import create_full_ui
+
+# Create the UI
+app = create_full_ui()
+
+# Launch the app
+if __name__ == "__main__":
+    app.launch(share=True, show_api=False)
+

requirements.txt

@@ -0,0 +1,75 @@
+aiofiles==23.2.1
+annotated-types==0.7.0
+anyio==4.8.0
+audioop-lts==0.2.1
+certifi==2025.1.31
+cffi==1.17.1
+charset-normalizer==3.4.1
+click==8.1.8
+colorama==0.4.6
+contourpy==1.3.1
+cycler==0.12.1
+fastapi==0.115.8
+ffmpy==0.5.0
+filelock==3.17.0
+fonttools==4.56.0
+fsspec==2025.2.0
+gevent==24.11.1
+gradio==5.15.0
+gradio_client==1.7.0
+greenlet==3.1.1
+h11==0.14.0
+httpcore==1.0.7
+httpx==0.28.1
+huggingface-hub==0.28.1
+idna==3.10
+Jinja2==3.1.5
+joblib==1.4.2
+kiwisolver==1.4.8
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.10.0
+mdurl==0.1.2
+mplfinance==0.12.10b0
+numpy==2.2.2
+nvdlib==0.7.9
+orjson==3.10.15
+packaging==24.2
+pandas==2.2.3
+pillow==11.1.0
+pycparser==2.22
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydub==0.25.1
+Pygments==2.19.1
+pyparsing==3.2.1
+python-dateutil==2.9.0.post0
+python-multipart==0.0.20
+pytz==2025.1
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.9.4
+ruff==0.9.6
+safehttpx==0.1.6
+scikit-learn==1.6.1
+scipy==1.15.1
+seaborn==0.13.2
+semantic-version==2.10.0
+setuptools==75.8.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+starlette==0.45.3
+threadpoolctl==3.5.0
+tomlkit==0.13.2
+tqdm==4.67.1
+typer==0.15.1
+typing_extensions==4.12.2
+tzdata==2025.1
+urllib3==2.3.0
+uvicorn==0.34.0
+websocket-client==1.8.0
+websockets==14.2
+zope.event==5.0
+zope.interface==7.2
+websocket-client

ui/FrontPage.py

@@ -0,0 +1,99 @@
+import gradio
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("./"))
+
+from source.api_calls.initiate_pipeline_call import initiate_pipeline_call
+from utils.process_steps import process_steps
+
+pipeline_steps = 21
+
+class FrontPage:
+    def __init__(self):
+        with gradio.Blocks(css=self.custom_css()) as self.page:
+            gradio.Markdown("# Code Repair with LLMs")
+            gradio.Markdown("Upload a Python, Java, C, or C++ file for processing")
+            
+            with gradio.Row():
+                file_input = gradio.File(label="Upload  Code File", file_types=[".py", ".java", ".c", ".cpp"])
+                language_input = gradio.Dropdown(choices=["Python", "Java", "C", "C++"], label="Language")
+            
+            file_content = gradio.Code(label="File Contents", language="python", interactive=False, elem_classes=["fixed-height"])
+            process_button = gradio.Button("Process")
+
+            with gradio.Row():
+                stage1 = gradio.Textbox(label="Fault Localization", value="Pending", interactive=False, elem_classes=["stage-box"])
+                stage2 = gradio.Textbox(label="Pattern Matching", value="Pending", interactive=False, elem_classes=["stage-box"])
+                stage3 = gradio.Textbox(label="Patch Generation", value="Pending", interactive=False, elem_classes=["stage-box"])
+                stage4 = gradio.Textbox(label="Patch Validation", value="Pending", interactive=False, elem_classes=["stage-box"])
+
+            output = gradio.Code(label="Processed Output", language="python", elem_classes=["fixed-height"])
+
+            pipeline_steps_input = gradio.Number(value=pipeline_steps, label="Pipeline Steps")
+
+            file_input.change(fn=self.display_file_content, inputs=[file_input], outputs=[file_content])
+            process_button.click(fn=self.initiate_pipeline, 
+                                 inputs=[file_input, language_input, pipeline_steps_input], outputs=[])
+
+    def custom_css(self):
+        return """
+        .fixed-height {
+            height: 300px !important;
+            overflow-y: auto !important;
+        }
+        .stage-box {
+            text-align: center !important;
+            font-weight: bold !important;
+        }
+        """
+
+    def display_file_content(self, file):
+        if file is None:
+            return "No file uploaded yet."
+        try:
+            if isinstance(file, str):
+                with open(file, 'r') as f:
+                    return f.read()
+            elif hasattr(file, 'name'):
+                return file.name
+            elif hasattr(file, 'read'):
+                return file.read().decode('utf-8')
+            else:
+                return str(file)
+        except Exception as e:
+            return f"An error occurred while reading the file: {str(e)}"
+
+
+    def initiate_pipeline(self, file_input, language_input, pipeline_steps):
+        if file_input is None or (isinstance(file_input, list) and len(file_input) == 0):
+            return "No file uploaded."
+
+        self.process_file(file_input, language_input)
+        
+        # If file_input is a list, take the first file
+        if isinstance(file_input, list):
+            file_input = file_input[0]
+            
+        print(f"Processing file: {file_input}")
+        return initiate_pipeline_call(file_input, pipeline_steps)
+    
+    def process_file(self, file, language):
+        if file is None:
+            return "Skipped", "Skipped", "Skipped", "Skipped", "Please upload a file."
+        
+        try:
+            content = self.display_file_content(file)
+            stage1_result = "Complete"
+            stage2_result = "Complete"
+            stage3_result = "Complete"
+            stage4_result = "Complete"
+            processed_content = self.process_content(content, language)
+            return stage1_result, stage2_result, stage3_result, stage4_result, processed_content
+        except Exception as e:
+            return "Error", "Error", "Error", "Error", f"An error occurred: {str(e)}"
+
+    def process_content(self, content, language):
+        # This is where you would implement your actual processing logic
+        # For now, we'll just return the file content with a message
+        return f"Processing {language} code:\n\n{content}\n\nProcessing complete."

utils/convert_steps.py

@@ -0,0 +1,4 @@
+def convert_steps(selected_steps):
+    """Converts selected steps into a binary value (bitmask)."""
+    all_steps = ["Bug Finding", "Pattern Matching", "Patch Generation", "Patch Validation"]
+    return sum(1 << i for i, step in enumerate(all_steps) if step in selected_steps)

utils/interface_utils.py

@@ -0,0 +1,105 @@
+# event handlers
+
+import gradio as gr
+
+def custom_theme():
+    # Gradio UI styling
+    return gr.themes.Ocean(
+        primary_hue=gr.themes.Color(c100="#fef9c3", c200="#fef08a", c300="#fde047", c400="#facc15", c50="#fefce8", c500="#eab308", c600="#ca8a04", c700="#a16207", c800="#854d0e", c900="#713f12", c950="#BA9B37"),
+        secondary_hue="zinc",
+        radius_size="lg",
+    ).set(
+        background_fill_primary='*neutral_900',
+        background_fill_secondary='*neutral_700',
+        body_background_fill='*secondary_900',
+        body_text_color='*neutral_100',
+        body_text_color_subdued='*neutral_400',
+        border_color_accent='*neutral_900',
+        button_secondary_background_fill='linear-gradient(120deg, *secondary_900 0%, *primary_400 50%, *primary_700 100%)',
+        button_secondary_background_fill_hover='linear-gradient(120deg, *secondary_900 0%, *primary_400 50%, *primary_700 100%)',
+        checkbox_label_background_fill_selected="linear-gradient(120deg, *secondary_900 0%, *primary_400 50%, *primary_700 100%)",
+        code_background_fill='*neutral_950',
+        color_accent_soft='*primary_950',
+        input_background_fill='*neutral_700',
+        table_odd_background_fill='*neutral_700'
+    )
+
+def custom_css():
+    return """
+    .fixed-height 
+    {
+        height: 300px !important;
+        overflow-y: auto !important;
+    }
+
+    .stage-box 
+    {
+        text-align: center !important;
+        font-weight: bold !important;
+    }
+
+    .title 
+    { 
+        text-align: center; 
+        font-size: 2rem; 
+        font-weight: bold; 
+    }
+
+    .subtitle 
+    { 
+        text-align: center;
+        font-size: 1.2rem;
+    }
+
+    .big-checkboxes label 
+    { 
+        font-size: 1.5rem;
+        display: block; 
+        text-align: center;
+    }
+
+    .big-button 
+    { 
+        font-size: 1.5rem; 
+        padding: 10px 20px; 
+        display: block; 
+        margin: 0 auto; 
+    }
+
+    """
+
+def display_file_content(file):
+    if file is None:
+        return "No file uploaded yet."
+    try:
+        if isinstance(file, str):
+            with open(file, 'r') as f:
+                return f.read()
+        elif hasattr(file, 'name'):
+            return file.name
+        elif hasattr(file, 'read'):
+            return file.read().decode('utf-8')
+        else:
+            return str(file)
+    except Exception as e:
+        return f"An error occurred while reading the file: {str(e)}"
+
+def process_file(file, language):
+    if file is None:
+        return "Skipped", "Skipped", "Skipped", "Skipped", "Please upload a file."
+    
+    try:
+        content = display_file_content(file)
+        stage1_result = "Complete"
+        stage2_result = "Complete"
+        stage3_result = "Complete"
+        stage4_result = "Complete"
+        processed_content = process_content(content, language)
+        return stage1_result, stage2_result, stage3_result, stage4_result, processed_content
+    except Exception as e:
+        return "Error", "Error", "Error", "Error", f"An error occurred: {str(e)}"
+
+def process_content(content, language):
+    # This is where you would implement your actual processing logic
+    # For now, we'll just return the file content with a message
+    return f"Processing {language} code:\n\n{content}\n\nProcessing complete."

utils/process_steps.py

@@ -0,0 +1,3 @@
+
+def process_steps():
+    return

utils/storage.py

@@ -0,0 +1,11 @@
+import gradio as gr
+
+def save_jwt_to_session(response):
+    """Extracts session_token and session_id from API response and saves them in Gradio State."""
+    session_token = response.get("session_token")
+    session_id = response.get("session_id")
+
+    if session_token and session_id:
+        return gr.update(value=session_token), gr.update(value=session_id)
+    
+    return gr.update(), gr.update()