116 lines
3.3 KiB
Python
116 lines
3.3 KiB
Python
import os
|
|
import logging
|
|
import requests
|
|
import csv
|
|
import datetime
|
|
from dotenv import load_dotenv
|
|
from telegram import Update
|
|
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackContext
|
|
|
|
load_dotenv()
|
|
|
|
logging.basicConfig(
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
|
level=logging.INFO
|
|
)
|
|
|
|
TELEGRAM_TOKEN = os.getenv("TELEGRAM_TOKEN")
|
|
FOLDER_ID = os.getenv("FOLDER_ID")
|
|
OAUTH_TOKEN = os.getenv("OAUTH_TOKEN")
|
|
|
|
|
|
STATS_FILE = "bot_statistics.csv"
|
|
|
|
def get_iam_token():
|
|
response = requests.post(
|
|
'https://iam.api.cloud.yandex.net/iam/v1/tokens',
|
|
json={'yandexPassportOauthToken': OAUTH_TOKEN}
|
|
)
|
|
|
|
response.raise_for_status()
|
|
return response.json()['iamToken']
|
|
|
|
def init_stats_file():
|
|
if not os.path.exists(STATS_FILE):
|
|
with open(STATS_FILE, 'w', newline='') as file:
|
|
writer = csv.writer(file)
|
|
writer.writerow(['timestamp', 'user_id', 'username', 'status'])
|
|
logging.info(f"Statistics file created at {STATS_FILE}")
|
|
|
|
def save_statistics(user_id, username, type="success"):
|
|
with open(STATS_FILE, 'a', newline='') as file:
|
|
writer = csv.writer(file)
|
|
timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
writer.writerow([timestamp, user_id, username, type])
|
|
|
|
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
|
"""Send a welcome message when the command /start is issued."""
|
|
user = update.effective_user
|
|
await update.message.reply_html(
|
|
f"Hi {user.mention_html()}! I'm your AI assistant powered by YandexGPT. Ask me anything!"
|
|
)
|
|
|
|
save_statistics(
|
|
user_id=user.id,
|
|
username=user.username or "unknown",
|
|
type="command_start"
|
|
)
|
|
|
|
async def process_message(update: Update, context: CallbackContext) -> None:
|
|
user_text = update.message.text
|
|
|
|
iam_token = get_iam_token()
|
|
|
|
data = {}
|
|
data["modelUri"] = f"gpt://{FOLDER_ID}/yandexgpt"
|
|
data["completionOptions"] = {"temperature": 0.3, "maxTokens": 1000}
|
|
data["messages"] = [
|
|
{"role": "user", "text": f"{user_text}"},
|
|
]
|
|
|
|
URL = "https://llm.api.cloud.yandex.net/foundationModels/v1/completion"
|
|
response = requests.post(
|
|
URL,
|
|
headers={
|
|
"Accept": "application/json",
|
|
"Authorization": f"Bearer {iam_token}"
|
|
},
|
|
json=data,
|
|
).json()
|
|
|
|
print(response)
|
|
|
|
save_statistics(
|
|
user_id=update.message.from_user.id,
|
|
username=update.effective_user.username,
|
|
type="Message"
|
|
)
|
|
|
|
answer = response.get('result', {})\
|
|
.get('alternatives', [{}])[0]\
|
|
.get('message', {})\
|
|
.get('text', {})
|
|
|
|
await update.message.reply_text(answer)
|
|
|
|
def main() -> None:
|
|
"""Start the bot."""
|
|
init_stats_file()
|
|
|
|
if not TELEGRAM_TOKEN:
|
|
print("No telegram token.")
|
|
return
|
|
|
|
print(TELEGRAM_TOKEN)
|
|
application = Application.builder().token(TELEGRAM_TOKEN).build()
|
|
|
|
application.add_handler(CommandHandler("start", start))
|
|
|
|
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, process_message))
|
|
|
|
application.run_polling()
|
|
logging.info("Bot started")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|