AIdea-Server / src /categorization /categorizer.py
Ali Hashhash
feat: enforce strict 7-category system and scaffold recommendations API endpoint
1ab78e2
from google import genai
from typing import Optional
from src.utils.logger import setup_logger
from src.utils.config import settings
logger = setup_logger(__name__)
class CategorizationService:
"""
Service for automatically categorizing notes based on their content using Gemini AI.
"""
def __init__(self, api_key: Optional[str] = None):
self.api_key = api_key or settings.gemini_api_key_1
# Use the newer google-genai client
self.client = genai.Client(api_key=self.api_key)
self.model_id = "gemini-1.5-flash-latest"
async def categorize_text(self, text: str) -> str:
"""
Categorize a piece of text into a single word or short phrase category.
"""
if not text or len(text) < 10:
return "Uncategorized"
prompt = (
"Analyze the following text and categorize it into exactly ONE of the following: "
"['Technology & AI', 'Business & Finance', 'Education & Science', 'Productivity & Self-Growth', 'News & Politics', 'Entertainment & Lifestyle', 'Health & Sports']. "
"Return ONLY the category string with no additional text or punctuation.\n\n"
f"Text: {text[:2000]}\n\n"
"Category:"
)
try:
# Using the newer google-genai syntax
response = self.client.models.generate_content(
model=self.model_id, contents=prompt
)
category = response.text.strip()
# Basic validation/cleanup
valid_categories = [
"Technology & AI",
"Business & Finance",
"Education & Science",
"Productivity & Self-Growth",
"News & Politics",
"Entertainment & Lifestyle",
"Health & Sports"
]
if category not in valid_categories:
return "Uncategorized"
return category
except Exception as e:
logger.error(f"Categorization failed: {e}")
return "Uncategorized"