File size: 741 Bytes
0eee603
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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