classify.py 1.5 KB

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