robosage / core_creator /intent_parser.py
mgbam's picture
Update core_creator/intent_parser.py
915b06b verified
raw
history blame
1.75 kB
# intent_parser.py - Smart fallback to Gemini or OpenAI for intent classification
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()
# Example
if __name__ == "__main__":
idea = "Build a robot that helps kids focus while doing homework."
print("Predicted Intent:", classify_robot_idea(idea))