classify.py 1.1 KB

1234567891011121314151617181920212223242526272829
  1. from nio.events.room_events import RoomMessageText
  2. from nio.rooms import MatrixRoom
  3. async def command_classify(room: MatrixRoom, event: RoomMessageText, bot):
  4. prompt = " ".join(event.body.split()[2:])
  5. if prompt:
  6. bot.logger.log("Classifying message...")
  7. try:
  8. response, tokens_used = await bot.classification_api.classify_message(prompt, user=room.room_id)
  9. except Exception as e:
  10. bot.logger.log(f"Error classifying message: {e}", "error")
  11. await bot.send_message(room, "Sorry, I couldn't classify the message. Please try again later.", True)
  12. return
  13. message = f"The message you provided seems to be of type: {response['type']}."
  14. if not prompt == response["prompt"]:
  15. message += f"\n\nPrompt: {response['prompt']}."
  16. await bot.send_message(room, message, True)
  17. bot.log_api_usage(event, room, f"{bot.classification_api.api_code}-{bot.classification_api.classification_api}", tokens_used)
  18. return
  19. await bot.send_message(room, "You need to provide a prompt.", True)