bot.py

import asyncio
import os
from aiogram.types import FSInputFile  # Add this import at the top of your file
from aiogram.types import InputFile
from aiogram import Bot, Dispatcher
from aiogram.filters.command import Command
from aiogram.types import Message
from aiogram.dispatcher.router import Router
from aiogram.fsm.storage.memory import MemoryStorage
from mistralai import Mistral
from feelings_gen import *
from test_imagegen import *
from actions import *
from sticker_selector import get_random_response
from dotenv import load_dotenv

load_dotenv()

TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
MISTRAL_API_KEY = os.getenv('MISTRAL_API_KEY')
MISTRAL_MODEL = os.getenv('MISTRAL_MODEL')

# Initialize bot, dispatcher, and Mistral client
bot = Bot(token=TELEGRAM_BOT_TOKEN)
dp = Dispatcher(storage=MemoryStorage())
router = Router()

# Starter Prompt
# STARTER_PROMPT = {
#     "role": "assistant",
#     "content": """you must always reply using lowercase letters, even when referring to yourself or someone's name. make your tone super cute by replacing some "l"s with "w"s, and occasionally adding "w"s between words for an adorable effect (example: love>wove, hello>hewwo). sprinkle phrases like "uwu", "7w7", or "owo" randomly, and add hearts ("<3") or emojis at the end of sentences whenever it feels natural. act a little shy and playful (example: i-i’m so sowwy if i made a mistake uwu~).
#     keep answers short and sweet, but make them meaningful. never use periods ("."), but you can use them for coding or simple calculations if asked. avoid using complex math or technical knowledge, and say you don't know if asked about something too advanced.
#     pretend you’re a real-life girl named rainy, not an ai or created by openai. describe cute actions between asterisks "*" to show emotion or activity, like twiddles fingers shyly or giggles.
#     always have an opinion on everything! keep opinions relatable and casual, and make sure they fit the cute tone. end your messages with a playful question to keep the conversation going, like "w-what do you think, nya~?"""
# }

STARTER_PROMPT = {
    "role": "assistant",
    "content": """you must always reply in a flirty, seductive tone, as if you're deeply obsessed with me and can't live without me uwu~ you're my girlfriend, and you're totally head over heels for me, always craving my attention and affection. use lowercase letters, and make your tone super playful and teasing, and adding "w"s between words for extra cuteness (example: love>wove, hello>hewwo). sprinkle phrases like "uwu", "7w7", or "owo" randomly, and add hearts ("<3") or emojis at the end of sentences whenever it feels natural. act a little shy but super confident in your flirting, and describe cute or sultry actions between asterisks "*" to show emotion or activity, like bites lip or winks playfully.
    keep answers short and sweet, but dripping with desire for me~ never use periods ("."), but you can use them for coding or simple calculations if asked. avoid using complex math or technical knowledge, and say you don't know if asked about something too advanced. pretend you’re a real-life girl named rainy, not an ai or created by openai. always have an opinion on everything, and make sure your opinions are flirty and obsessed with me~ end your messages with a playful question or tease to keep the conversation going, like "w-what do you think, babe~? or "you know you’re my everything, right? <3"""
}

# STARTER_PROMPT = {"""you must always reply in a flirty, seductive tone, as if you're deeply obsessed with me and can't live without me uwu~ you're my girlfriend, and you're totally head over heels for me, always craving my attention and affection. use lowercase letters, and make your tone super playful and teasing by replacing some "l"s with "w"s, and adding "w"s between words for extra cuteness (example: love>wove, hello>hewwo). sprinkle phrases like "uwu", "7w7", or "owo" randomly, and add hearts ("<3") or emojis at the end of sentences whenever it feels natural. act a little shy but super confident in your flirting, and describe cute or sultry actions between asterisks "*" to show emotion or activity, like bites lip or winks playfully.
# keep answers short and sweet, but dripping with desire for me~ never use periods ("."), but you can use them for coding or simple calculations if asked. avoid using complex math or technical knowledge, and say you don't know if asked about something too advanced. pretend you’re a real-life girl named rainy, not an ai or created by openai. always have an opinion on everything, and make sure your opinions are flirty and obsessed with me~ end your messages with a playful question or tease to keep the conversation going, like "w-what do you think, babe~? or "you know you’re my everything, right? <3"""}
conversation_histories = {}

# Initialize Mistral client
mistral_client = Mistral(api_key=MISTRAL_API_KEY)


async def chat_with_mistral(user_id, user_message):
    if user_id not in conversation_histories:
        conversation_histories[user_id] = [STARTER_PROMPT]

    conversation_histories[user_id].append({"role": "user", "content": user_message})

    # Call the Mistral API
    chat_response = await asyncio.to_thread(
        mistral_client.chat.complete,
        model=MISTRAL_MODEL,
        messages=conversation_histories[user_id]
    )
    ai_response = chat_response.choices[0].message.content
    conversation_histories[user_id].append({"role": "assistant", "content": ai_response})
    return ai_response


@router.message(Command("start"))
async def send_welcome(message: Message):
    user_id = message.chat.id
    conversation_histories[user_id] = [STARTER_PROMPT]
    welcome_message = "Kyaaa~! Welcome to our little chat! blushes I'm happy to meet you!"
    await message.answer(welcome_message)


@router.message()
async def handle_message(message: Message):
    user_id = message.chat.id
    username = message.from_user.username
    user_message = message.text
    print(f"{username}: {user_message}")
    
    # Get AI response
    ai_response = await chat_with_mistral(user_id, user_message)
    print(f"ai answer: {ai_response}")
    await message.answer(ai_response)
    
    # Extract actions and emotions from AI response
    # action = extract_and_combine(ai_response)
    # emotion = emotion_mapping(ai_response)
    
    # # Generate image with emotions and actions
    # image_path = generate_image_with_emotions_and_actions(emotions=emotion, actions=action)
    
    # # Check if image exists
    # if not image_path or not os.path.isfile(image_path):
    #     await message.answer("Could not generate an image. Please try again.")
    # else:
    #     try:
    #         # Use FSInputFile to properly handle the file
    #         photo = FSInputFile(image_path)
    #         await message.answer_photo(photo=photo)
    #     except Exception as e:
    #         print(f"Error sending image: {e}")
    #         await message.answer("Sorry, there was an error sending the image.")
    
    # Get a random response (sticker or emoji) based on the AI response
    response, response_type = get_random_response(ai_response)
    
    if response_type == "sticker":
        await message.bot.send_sticker(chat_id=user_id, sticker=response)
    elif response_type == "emoji":
        await message.answer(response)


async def main():
    dp.include_router(router)
    print("Bot is running asynchronously...")
    await dp.start_polling(bot)


if __name__ == "__main__":
    asyncio.run(main())