import os
import gradio as gr
import openai

# Initialize OpenAI API key
openai.api_key = os.environ.get("OPENAI_API_KEY")

def predict(message, history):
    """
    Function to handle chat interactions with streaming responses
    """
    # Convert history to the format expected by OpenAI API
    messages = []
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        if h[1]:
            messages.append({"role": "assistant", "content": h[1]})
    
    # Add the current message
    messages.append({"role": "user", "content": message})
    
    # Stream the response using the older API style
    response = openai.ChatCompletion.create(
        model="ft:gpt-4o-mini-2024-07-18:ignitarium-technology-solutions:test-case:BAgDvOiS", 
        messages=messages,
        stream=True
    )
    
    partial_response = ""
    for chunk in response:
        if 'content' in chunk['choices'][0].get('delta', {}):
            content = chunk['choices'][0]['delta']['content']
            partial_response += content
            yield partial_response

# Create the chat interface with streaming
demo = gr.ChatInterface(
    fn=predict,
    title="Chat with Fine-tuned GPT-4o-mini (Streaming)",
    description="Interact with my fine-tuned GPT-4o-mini model with streaming responses",
    examples=["""Design a complete set of test cases for validating SD card boot functionality organized by the following categories:

Storage capacity scenarios (full memory, minimal free space, varying card sizes)

Boot sequence verification (boot order priorities, multiple boot media present)

File system compatibility (FAT16/32, ext4, and other supported formats)

Hardware interaction (card insertion/removal, power cycling during boot)

Performance testing (boot time with varying card speeds/classes)

Error handling and recovery (corrupted boot sectors, partial writes)

Security aspects (if applicable)

For each test case, provide a clear title, preconditions, detailed steps, expected outcomes, and pass/fail criteria. Include both positive and negative test scenarios to ensure robust validation""", """Create a complete set of test cases for validating I2C boot functionality organized by the following categories:

Bus initialization and configuration (address assignment, speed configuration)

Boot sequence verification (proper execution of boot stages via I2C)

Protocol compliance (start/stop conditions, acknowledgment handling)

Error detection and recovery (noise immunity, timeout handling, retry mechanisms)

Performance testing (boot time at different I2C clock frequencies)

Multi-device scenarios (multiple I2C devices on the bus during boot)

Power management (low power boot options, power cycling effects)

Security aspects (if applicable)

For each test case, provide a descriptive title, preconditions, step-by-step procedures, expected outcomes, and clear pass/fail criteria. Include both positive test cases to verify correct functionality and negative test cases to validate proper error handling and system robustness."""]
)

# Launch the app
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", ssr_mode=False)