|
import openai
|
|
from typing import Dict, Any, Optional
|
|
|
|
class AIProjectAssistant:
|
|
"""Wrapper class for OpenAI API interactions focused on Data Science and AI projects."""
|
|
|
|
def __init__(self, api_key: str):
|
|
"""Initialize the AI assistant with OpenAI API key."""
|
|
openai.api_key = api_key
|
|
self.default_model = "gpt-4"
|
|
self.default_temperature = 0.7
|
|
|
|
def send_prompt(self,
|
|
prompt: str,
|
|
model: Optional[str] = None,
|
|
temperature: Optional[float] = None) -> Dict[str, Any]:
|
|
"""
|
|
Send a prompt to OpenAI API and get the response.
|
|
|
|
Args:
|
|
prompt (str): The user's input prompt
|
|
model (str, optional): OpenAI model to use. Defaults to gpt-4
|
|
temperature (float, optional): Response creativity (0-1). Defaults to 0.7
|
|
|
|
Returns:
|
|
Dict[str, Any]: OpenAI API response
|
|
"""
|
|
try:
|
|
response = openai.ChatCompletion.create(
|
|
model=model or self.default_model,
|
|
messages=[
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
temperature=temperature or self.default_temperature
|
|
)
|
|
return {
|
|
"success": True,
|
|
"content": response.choices[0].message.content,
|
|
"tokens_used": response.usage.total_tokens
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"success": False,
|
|
"error": str(e)
|
|
}
|
|
|
|
def brainstorm_project(self, topic: str) -> Dict[str, Any]:
|
|
"""
|
|
Generate AI project ideas based on a topic.
|
|
|
|
Args:
|
|
topic (str): The topic or field of interest
|
|
|
|
Returns:
|
|
Dict[str, Any]: Project suggestions and implementation details
|
|
"""
|
|
prompt = f"""
|
|
Help me brainstorm an AI/Data Science project related to {topic}. Please provide:
|
|
1. Project title
|
|
2. Problem statement
|
|
3. Suggested approach
|
|
4. Required technologies/libraries
|
|
5. Potential challenges
|
|
6. Expected outcomes
|
|
"""
|
|
return self.send_prompt(prompt)
|
|
|
|
def get_code_suggestion(self, description: str) -> Dict[str, Any]:
|
|
"""
|
|
Get code suggestions for implementing specific functionality.
|
|
|
|
Args:
|
|
description (str): Description of the desired functionality
|
|
|
|
Returns:
|
|
Dict[str, Any]: Code suggestions and explanations
|
|
"""
|
|
prompt = f"""
|
|
Please provide Python code suggestions for the following functionality:
|
|
{description}
|
|
|
|
Include:
|
|
1. Code implementation
|
|
2. Required imports
|
|
3. Brief explanation of the approach
|
|
4. Any potential improvements
|
|
"""
|
|
return self.send_prompt(prompt) |