import openai
import os
from dotenv import load_dotenv

load_dotenv()

openai.api_key = os.getenv("OPENAI_API_KEY")

def code_agent(prompt, context=None):
    """
    Generate code based on the prompt and optional context from RAG.
    """
    if context:
        prompt = f"Use the following context to help with the task:\n{context}\n\nTask: {prompt}"

    response = openai.ChatCompletion.create(
        model="gpt-4",  # or "gpt-3.5-turbo" if you prefer
        messages=[
            {"role": "system", "content": "You are a helpful coding assistant. Write code in Python."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.5,
        max_tokens=500
    )
    return response.choices[0].message.content