feat: add local API support for OpenAI compatible APIs

Signed-off-by: salmonik <raisersoft@proton.me>

1 files changed, 35 insertions, 40 deletions

diff --git a/cogs/ai.py b/cogs/ai.py
index ef2ba7c..dbdb02e 100644
--- a/cogs/ai.py
+++ b/cogs/ai.py
@@ -7,12 +7,14 @@ import os
 import discord
 from discord.ext import commands
 import logging
-import httpx
 from gtts import gTTS
 import io
+from cogs.utils.llm import llm_completion
 
 logger = logging.getLogger(__name__)
 
+_DEFAULT_SENATOR_MODEL = "cognitivecomputations/dolphin-mistral-24b-venice-edition:free"
+
 class AI(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
@@ -29,48 +31,41 @@ class AI(commands.Cog):
     ):
         await ctx.defer()
 
+        model = (
+            os.getenv("LLM_MODEL")
+            or os.getenv("OPENROUTER_MODEL")
+            or _DEFAULT_SENATOR_MODEL
+        )
+        messages = [
+            {
+                "role": "system",
+                "content": "You are Senator Steaven Sperma Armstrong from Metal Gear Rising Vengeance. \
+                You are racist and you hate other minorities. Respond in plain text, no markdown tricks. \
+                If the message is in Polish, respond in Polish. Respond in 100 words."
+            },
+            {
+                "role": "user",
+                "content": text
+            }
+        ]
 
-        url = "https://openrouter.ai/api/v1/chat/completions"
-        headers = {
-                "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
-                "Content-Type": "application/json",
-         }
-        payload = {
-             "model": "cognitivecomputations/dolphin-mistral-24b-venice-edition:free",
-             "messages": [
-                 {
-                     "role": "system",
-                     "content": "You are Senator Steaven Sperma Armstrong from Metal Gear Rising Vengeance. \
-                     You are racist and you hate other minorities. Respond in plain text, no markdown tricks. \
-                     If the message is in Polish, respond in Polish. Respond in 100 words."
-                 },
-                 {
-                     "role": "user",
-                     "content": text
-                 }
-             ],
+        response_text = await llm_completion(messages, model=model)
 
-         }
-        async with httpx.AsyncClient() as client:
-            response = await client.post(url=url, headers=headers, json=payload)
-            response_data = response.json()
-            logger.info(response_data)
-            if not response_data["choices"]:
-                await ctx.send_followup("bot siem zesroł")
-            else:
-                response_text = response_data["choices"][0]["message"]["content"]
-                if tts:
-                    tts_obj = gTTS(text=response_text, lang='pl')
-                    tts_buffer = io.BytesIO()
-                    tts_obj.write_to_fp(tts_buffer)
-                    tts_buffer.seek(0)
-                    from cogs.voice import get_voice_manager
-                    vm = get_voice_manager(self.bot)
-                    channel = ctx.author.voice.channel
-                    await vm.play_source(channel, discord.FFmpegPCMAudio(tts_buffer), interrupt=False)
+        if response_text is None:
+            await ctx.send_followup("bot siem zesroł")
+            return
 
-                else:
-                    await ctx.send_followup(response_text)
+        if tts:
+            tts_obj = gTTS(text=response_text, lang='pl')
+            tts_buffer = io.BytesIO()
+            tts_obj.write_to_fp(tts_buffer)
+            tts_buffer.seek(0)
+            from cogs.voice import get_voice_manager
+            vm = get_voice_manager(self.bot)
+            channel = ctx.author.voice.channel
+            await vm.play_source(channel, discord.FFmpegPCMAudio(tts_buffer), interrupt=False)
+        else:
+            await ctx.send_followup(response_text)