Add Ollama bot

catch-all/06_bots_telegram/09_ollama_bot/bot/func/interactions.py

@@ -0,0 +1,199 @@
+# >> interactions
+import logging
+import os
+import aiohttp
+import json
+
+from aiogram import types
+from aiohttp import ClientTimeout
+from asyncio import Lock
+from functools import wraps
+from dotenv import load_dotenv
+
+
+load_dotenv('.env')
+token = os.getenv("TOKEN")
+
+allowed_ids = list(map(int, os.getenv("USER_IDS", "").split(",")))
+admin_ids = list(map(int, os.getenv("ADMIN_IDS", "").split(",")))
+
+ollama_base_url = os.getenv("OLLAMA_BASE_URL")
+ollama_port = os.getenv("OLLAMA_PORT", "11434")
+
+log_level_str = os.getenv("LOG_LEVEL", "INFO")
+
+allow_all_users_in_groups = bool(
+    int(os.getenv("ALLOW_ALL_USERS_IN_GROUPS", "0")))
+
+log_levels = list(logging._levelToName.values())
+
+timeout = os.getenv("TIMEOUT", "3000")
+
+if log_level_str not in log_levels:
+
+    log_level = logging.DEBUG
+
+else:
+
+    log_level = logging.getLevelName(log_level_str)
+
+logging.basicConfig(level=log_level)
+
+
+async def model_list():
+
+    async with aiohttp.ClientSession() as session:
+
+        url = f"http://{ollama_base_url}:{ollama_port}/api/tags"
+
+        async with session.get(url) as response:
+
+            if response.status == 200:
+
+                data = await response.json()
+                return data["models"]
+
+            else:
+
+                return []
+
+
+async def generate(payload: dict, modelname: str, prompt: str):
+
+    client_timeout = ClientTimeout(total=int(timeout))
+
+    async with aiohttp.ClientSession(timeout=client_timeout) as session:
+
+        url = f"http://{ollama_base_url}:{ollama_port}/api/chat"
+
+        try:
+
+            async with session.post(url, json=payload) as response:
+
+                if response.status != 200:
+
+                    raise aiohttp.ClientResponseError(
+
+                        status=response.status, message=response.reason
+
+                    )
+
+                buffer = b""
+
+                async for chunk in response.content.iter_any():
+
+                    buffer += chunk
+
+                    while b"\n" in buffer:
+
+                        line, buffer = buffer.split(b"\n", 1)
+                        line = line.strip()
+
+                        if line:
+
+                            yield json.loads(line)
+
+        except aiohttp.ClientError as e:
+
+            print(f"Error during request: {e}")
+
+
+def perms_allowed(func):
+
+    @wraps(func)
+    async def wrapper(message: types.Message = None, query: types.CallbackQuery = None):
+
+        user_id = message.from_user.id if message else query.from_user.id
+
+        if user_id in admin_ids or user_id in allowed_ids:
+
+            if message:
+
+                return await func(message)
+
+            elif query:
+
+                return await func(query=query)
+
+        else:
+
+            if message:
+
+                if message and message.chat.type in ["supergroup", "group"]:
+
+                    if allow_all_users_in_groups:
+
+                        return await func(message)
+
+                    return
+
+                await message.answer("Access Denied")
+
+            elif query:
+
+                if message and message.chat.type in ["supergroup", "group"]:
+
+                    return
+
+                await query.answer("Access Denied")
+
+    return wrapper
+
+
+def perms_admins(func):
+
+    @wraps(func)
+    async def wrapper(message: types.Message = None, query: types.CallbackQuery = None):
+
+        user_id = message.from_user.id if message else query.from_user.id
+
+        if user_id in admin_ids:
+
+            if message:
+
+                return await func(message)
+
+            elif query:
+
+                return await func(query=query)
+
+        else:
+
+            if message:
+
+                if message and message.chat.type in ["supergroup", "group"]:
+
+                    return
+
+                await message.answer("Access Denied")
+
+                logging.info(
+                    f"[MSG] {message.from_user.first_name} {
+                        message.from_user.last_name}({message.from_user.id}) is not allowed to use this bot."
+                )
+
+            elif query:
+
+                if message and message.chat.type in ["supergroup", "group"]:
+
+                    return
+
+                await query.answer("Access Denied")
+
+                logging.info(
+                    f"[QUERY] {message.from_user.first_name} {
+                        message.from_user.last_name}({message.from_user.id}) is not allowed to use this bot."
+                )
+
+    return wrapper
+
+
+class contextLock:
+
+    lock = Lock()
+
+    async def __aenter__(self):
+        await self.lock.acquire()
+
+    async def __aexit__(self, exc_type, exc_value, exc_traceback):
+        self.lock.release()