| | |
| |
|
| | 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)) |
| |
|