95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
import google.generativeai as genai
|
|
import os
|
|
|
|
def classify_email(api_key, email_content, categories):
|
|
"""
|
|
Classifies the given email content into one of the provided categories using the Gemini API.
|
|
|
|
Args:
|
|
api_key (str): The Gemini API key.
|
|
email_content (str): The content of the email (subject + body).
|
|
categories (list): A list of strings representing the possible categories.
|
|
|
|
Returns:
|
|
str: The chosen category, or None if classification fails.
|
|
"""
|
|
if not api_key or api_key == "YOUR_GEMINI_API_KEY":
|
|
print("! 경고: Gemini API 키가 설정되지 않았습니다. AI 분류를 건너뜁니다.")
|
|
return None
|
|
|
|
try:
|
|
genai.configure(api_key=api_key)
|
|
except Exception as e:
|
|
print(f"Gemini API 설정 중 오류 발생: {e}")
|
|
return None
|
|
|
|
# For safety, let's use a model that's less likely to refuse classification
|
|
model = genai.GenerativeModel('gemini-1.5-flash')
|
|
|
|
# Construct the prompt
|
|
prompt = f"""
|
|
Analyze the following email content and classify it into one of the following categories.
|
|
Your response must be ONLY ONE of the category names provided. Do not add any extra text, explanation, or punctuation.
|
|
|
|
Categories: {', '.join(categories)}
|
|
|
|
Email Content:
|
|
---
|
|
{email_content}
|
|
---
|
|
|
|
Category:
|
|
"""
|
|
|
|
try:
|
|
response = model.generate_content(prompt)
|
|
# The response text should be one of the categories
|
|
result_category = response.text.strip()
|
|
|
|
# Validate that the model returned a valid category
|
|
if result_category in categories:
|
|
print(f"-> AI 분류 결과: '{result_category}'")
|
|
return result_category
|
|
else:
|
|
print(f"! 경고: AI가 유효하지 않은 카테고리 '{result_category}'를 반환했습니다.")
|
|
# Fallback or simply return None
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"Gemini API 호출 중 오류 발생: {e}")
|
|
return None
|
|
|
|
if __name__ == '__main__':
|
|
# --- 사용 예시 ---
|
|
# 1. 아래 YOUR_API_KEY를 실제 키로 변경하세요.
|
|
# 2. 터미널에서 `pip install google-generativeai` 를 실행하세요.
|
|
# 3. `python ai_classifier.py` 명령으로 테스트하세요.
|
|
|
|
test_api_key = os.environ.get("GEMINI_API_KEY", "YOUR_GEMINI_API_KEY") # 실제 키로 교체 필요
|
|
|
|
test_email = """
|
|
Subject: Invoice INV-2024-001 from Office Supplies Co.
|
|
|
|
Dear Customer,
|
|
|
|
Please find attached your invoice for the recent purchase of office supplies.
|
|
The total amount due is $125.50.
|
|
|
|
Thank you for your business.
|
|
|
|
Sincerely,
|
|
Office Supplies Co.
|
|
"""
|
|
|
|
test_categories = ["청구서", "프로젝트 업데이트", "광고", "개인 용무"]
|
|
|
|
if test_api_key != "YOUR_GEMINI_API_KEY":
|
|
classified_category = classify_email(test_api_key, test_email, test_categories)
|
|
|
|
if classified_category:
|
|
print(f"\n최종 분류: {classified_category}")
|
|
else:
|
|
print("\n분류에 실패했습니다.")
|
|
else:
|
|
print("\n테스트를 위해 'YOUR_GEMINI_API_KEY'를 실제 API 키로 바꿔주세요.")
|