|
|
|
|
|
import os |
|
import json |
|
|
|
USE_GEMINI = bool(os.getenv("GEMINI_API_KEY")) |
|
USE_OPENAI = bool(os.getenv("OPENAI_API_KEY")) |
|
|
|
INTENT_CATEGORIES = [ |
|
"educational", |
|
"assistive", |
|
"entertainment", |
|
"industrial", |
|
"home automation", |
|
"healthcare", |
|
"retail", |
|
"creative" |
|
] |
|
|
|
if USE_GEMINI: |
|
import google.generativeai as genai |
|
genai.configure(api_key=os.getenv("GEMINI_API_KEY")) |
|
gemini_model = genai.GenerativeModel(model_name="models/gemini-1.5-pro-latest") |
|
|
|
elif USE_OPENAI: |
|
from openai import OpenAI |
|
openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) |
|
else: |
|
raise EnvironmentError("No valid API key set. Please set GEMINI_API_KEY or OPENAI_API_KEY.") |
|
|
|
def classify_robot_idea(user_input: str) -> str: |
|
prompt = f""" |
|
Classify this user idea into one of the following categories: |
|
{', '.join(INTENT_CATEGORIES)}. |
|
Only return the category word. If none fits, return 'creative'. |
|
|
|
Idea: {user_input} |
|
Category: |
|
""" |
|
|
|
if USE_GEMINI: |
|
response = gemini_model.generate_content(prompt) |
|
return response.text.strip().lower() |
|
|
|
elif USE_OPENAI: |
|
response = openai_client.chat.completions.create( |
|
model="gpt-4o", |
|
messages=[ |
|
{"role": "system", "content": "You are a classification AI for robotics ideas."}, |
|
{"role": "user", "content": prompt} |
|
], |
|
temperature=0 |
|
) |
|
return response.choices[0].message.content.strip().lower() |
|
|
|
|
|
if __name__ == "__main__": |
|
idea = "Build a robot that helps kids focus while doing homework." |
|
print("Predicted Intent:", classify_robot_idea(idea)) |
|
|