import gradio as gr
import OpenAI

client = OpenAI(
    base_url="http://localhost:11434/v1/",
    api_key="ollama"  # Required but unused by Ollama
)

def chat_with_code(message, history):
    response = client.chat.completions.create(
        model="llama3.2",
        messages=[{"role": "user", "content": message}],
        temperature=0.7
    )
    
    full_response = response.choices[0].message.content
    code_block = None
    
    # Extract code blocks from response
    if "```python" in full_response:
        code = full_response.split("```python")[1].split("```")[0].strip()
        code_block = gr.Code(language="python", value=code)
    elif "```javascript" in full_response:
        code = full_response.split("```javascript")[1].split("```")[0].strip()
        code_block = gr.Code(language="javascript", value=code)
    
    return full_response, code_block

with gr.Blocks() as demo:
    code = gr.Code(render=False)
    gr.ChatInterface(
        fn=chat_with_code,
        additional_outputs=[code],
        type="messages"
    )

demo.launch()