Spaces:
Build error
Build error
File size: 3,193 Bytes
afc2b98 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import openai
import os
from typing import Dict, Any
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
class CoderAgent:
"""
Agent responsible for generating code based on user prompts.
This agent understands programming requirements and writes initial code.
"""
def __init__(self, model="gpt-3.5-turbo", temperature=0.7):
"""
Initialize the coder agent with specific model parameters.
Args:
model: Which OpenAI model to use
temperature: Controls creativity (0=deterministic, 1=creative)
"""
self.model = model
self.temperature = temperature
self.api_key = os.getenv("OPENAI_API_KEY")
if not self.api_key:
raise ValueError("OPENAI_API_KEY not found in environment variables")
openai.api_key = self.api_key
def generate_code(self, prompt: str, context: str = "") -> Dict[str, Any]:
"""
Generate code based on user prompt and optional context.
Args:
prompt: User's coding requirement (e.g., "Write a function to reverse a string")
context: Additional context from RAG or previous conversations
Returns:
Dictionary containing code and metadata
"""
try:
# Build the system message
system_message = """You are an expert Python programmer.
Write clean, efficient, and well-documented code.
Always include function signatures with type hints.
Provide only the code without explanations unless asked."""
# Build user message with context
user_message = prompt
if context:
user_message = f"Context: {context}\n\nTask: {prompt}"
# Call OpenAI API
response = openai.ChatCompletion.create(
model=self.model,
messages=[
{"role": "system", "content": system_message},
{"role": "user", "content": user_message}
],
temperature=self.temperature,
max_tokens=500
)
# Extract and clean the code
raw_code = response.choices[0].message.content
# Clean the response (remove markdown code blocks if present)
if "```python" in raw_code:
code = raw_code.split("```python")[1].split("```")[0].strip()
elif "```" in raw_code:
code = raw_code.split("```")[1].split("```")[0].strip()
else:
code = raw_code.strip()
return {
"status": "success",
"code": code,
"raw_response": raw_code,
"model_used": self.model,
"tokens_used": response.usage.total_tokens
}
except Exception as e:
return {
"status": "error",
"error": str(e),
"code": "",
"raw_response": ""
} |