aiprojectideas / ai_wrapper.py
decodingdatascience's picture
Upload 5 files
b6263d3 verified
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)