droppler/main.py at main · LanCoCafe/droppler · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
from os import getenv

from dotenv import load_dotenv
from flask import Flask, request, abort, redirect
from google import generativeai
from google.generativeai import GenerativeModel
from linebot.v3 import (
    WebhookHandler
)
from linebot.v3.exceptions import (
    InvalidSignatureError
)
from linebot.v3.messaging import (
    Configuration, ApiClient, MessagingApi
)
from linebot.v3.webhooks import (
    MessageEvent,
    TextMessageContent
)
from pymongo import MongoClient
from uvicorn import Config, Server

from src.process_messages import process_group_message, process_user_message

load_dotenv()

app = Flask(__name__)

configuration = Configuration(access_token=getenv("CHANNEL_ACCESS_TOKEN"))
handler = WebhookHandler(getenv("CHANNEL_SECRET"))

api_client = ApiClient(configuration)
line_bot_api = MessagingApi(api_client)

database = MongoClient(getenv("MONGODB_URI")).get_database("droppler")


def setup_gemini() -> GenerativeModel:
    generativeai.configure(api_key=getenv("GEMINI_API_KEY"))
    generation_config = {
        "temperature": 0.9,
        "top_p": 1,
        "top_k": 1,
        "max_output_tokens": 2048,
        "response_mime_type": "application/json",
    }

    safety_settings = [
        {
            "category": "HARM_CATEGORY_HARASSMENT",
            "threshold": "BLOCK_NONE"
        },
        {
            "category": "HARM_CATEGORY_HATE_SPEECH",
            "threshold": "BLOCK_NONE"
        },
        {
            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
            "threshold": "BLOCK_NONE"
        },
        {
            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
            "threshold": "BLOCK_NONE"
        },
    ]

    return GenerativeModel(
        model_name="gemini-1.5-flash",
        generation_config=generation_config,
        safety_settings=safety_settings
    )


model = setup_gemini()


@app.route("/", methods=['GET'])
def root():
    return redirect("https://wolf-yuan.dev", code=301)


@app.route("/callback", methods=['POST'])
def callback():
    signature = request.headers['X-Line-Signature']

    body = request.get_data(as_text=True)
    app.logger.info("Request body: " + body)

    try:
        handler.handle(body, signature)
    except InvalidSignatureError:
        app.logger.info("Invalid signature. Please check your channel access token/channel secret.")
        abort(400)

    return 'OK'


@handler.add(MessageEvent, message=TextMessageContent)
def handle_message(event: MessageEvent):
    if not event.source:
        return

    elif event.source.type == "user":
        process_user_message(line_bot_api, event)

    elif event.source.type == "group":
        process_group_message(database, line_bot_api, event, model)

    else:
        return


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=3002)