Merge pull request 'feat: add local API support for OpenAI compatible APIs' (#6) from dev into main

Reviewed-on: https://codeberg.org/salmonik/senator-sperma/pulls/6

5 files changed, 122 insertions, 136 deletions

diff --git a/.env.example b/.env.example
index 9623155..b7c7c61 100644
--- a/.env.example
+++ b/.env.example
@@ -11,5 +11,9 @@ COPYPARTY_URL_UPLOAD=http://localhost:8080/uploads
 COPYPARTY_PASSWORD=
 
 # Optional - AI commands
-OPENROUTER_API_KEY=
-OPENROUTER_MODEL=
+# LLM_API_URL=http://localhost:8080/v1/chat/completions   # llama.cpp server (use http://, no TLS)
+# LLM_API_KEY=                                            # not needed for llama.cpp
+# LLM_API_VERIFY=true                                     # SSL verify (false for self-signed certs)
+# LLM_MODEL=                                              # model name for llama.cpp
+OPENROUTER_API_KEY=   # deprecated, use LLM_API_KEY instead
+OPENROUTER_MODEL=     # deprecated, use LLM_MODEL instead
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)
 
 
 
diff --git a/cogs/assistant.py b/cogs/assistant.py
index fed7825..ffb664e 100644
--- a/cogs/assistant.py
+++ b/cogs/assistant.py
@@ -5,9 +5,9 @@
 
 import discord
 from discord.ext import commands
-import httpx
 import logging
 import os
+from cogs.utils.llm import llm_completion
 
 logger = logging.getLogger(__name__)
 
@@ -32,24 +32,20 @@ class Jarvis(commands.Cog):
 
     @staticmethod
     async def assistant(message):
-        headers = {
-            "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
-            "Content-Type": "application/json"
-        }
-
         chain = await Jarvis._build_message_chain(message)
         messages = Jarvis._build_messages(message, chain)
-        payload = Jarvis._build_payload(messages)
 
-        response_data = await Jarvis._safe_api_call(payload, headers)
-        if response_data is None:
+        extra = None
+        api_url = os.getenv("LLM_API_URL", "https://openrouter.ai/api/v1/chat/completions")
+        if "openrouter" in api_url.lower():
+            extra = {"plugins": [{"id": "response-healing"}]}
+
+        response_text = await llm_completion(messages, extra=extra)
+        if response_text is None:
             await message.reply("Przepraszam, nie mogę teraz odpowiedzieć. Spróbuj ponownie później.")
             return
 
-        response_message = Jarvis._parse_response(response_data)
-
-        channel = message.channel
-        await message.reply(response_message)
+        await message.reply(response_text)
 
     @staticmethod
     async def _build_message_chain(message):
@@ -93,38 +89,7 @@ class Jarvis(commands.Cog):
 
         return messages
 
-    @staticmethod
-    def _build_payload(messages):
-        return {
-            "model": os.getenv("OPENROUTER_MODEL"),
-            "messages": messages,
-            "plugins": [{"id": "response-healing"}]
-        }
-
-    @staticmethod
-    async def _make_api_request(payload, headers):
-        async with httpx.AsyncClient(timeout=30.0) as client:
-            url = "https://openrouter.ai/api/v1/chat/completions"
-            response = await client.post(url=url, headers=headers, json=payload)
-            response_data = response.json()
-            logger.info(response_data)
-            return response_data
-
-    @staticmethod
-    async def _safe_api_call(payload, headers):
-        try:
-            return await Jarvis._make_api_request(payload, headers)
-        except httpx.HTTPError as e:
-            logger.error(f"API request failed: {e}")
-            return None
-
-    @staticmethod
-    def _parse_response(response_data):
-        content = response_data['choices'][0]['message']['content']
-
-        logger.info(content)
 
-        return content
 
 def setup(bot):
     bot.add_cog(Jarvis(bot))
diff --git a/cogs/utils/llm.py b/cogs/utils/llm.py
new file mode 100644
index 0000000..2b05b12
--- /dev/null
+++ b/cogs/utils/llm.py
@@ -0,0 +1,62 @@
+# This file is part of Senator Sperma
+# Licensed under the GPL-3.0 License
+# Copyright (C) 2026 Salmonik @ salmonik@tuta.io
+# https://codeberg.org/salmonik/senator-sperma
+
+import os
+import httpx
+import logging
+
+logger = logging.getLogger(__name__)
+
+BACKWARD_COMPAT_MSG = (
+    "OPENROUTER_API_KEY/OPENROUTER_MODEL are deprecated. "
+    "Use LLM_API_KEY, LLM_MODEL instead."
+)
+
+
+def _get_config():
+    api_url = os.getenv("LLM_API_URL", "https://openrouter.ai/api/v1/chat/completions")
+    api_key = os.getenv("LLM_API_KEY") or os.getenv("OPENROUTER_API_KEY") or ""
+    model = os.getenv("LLM_MODEL") or os.getenv("OPENROUTER_MODEL") or ""
+
+    if os.getenv("OPENROUTER_API_KEY") and not os.getenv("LLM_API_KEY"):
+        logger.warning(BACKWARD_COMPAT_MSG)
+    if os.getenv("OPENROUTER_MODEL") and not os.getenv("LLM_MODEL"):
+        logger.warning(BACKWARD_COMPAT_MSG)
+
+    return api_url, api_key, model
+
+
+async def llm_completion(
+    messages: list[dict],
+    model: str | None = None,
+    api_url: str | None = None,
+    api_key: str | None = None,
+    extra: dict | None = None,
+    timeout: float = 180.0,
+) -> str | None:
+    url = api_url or os.getenv("LLM_API_URL", "https://openrouter.ai/api/v1/chat/completions")
+    model = model or os.getenv("LLM_MODEL") or os.getenv("OPENROUTER_MODEL") or ""
+    api_key = api_key or os.getenv("LLM_API_KEY") or os.getenv("OPENROUTER_API_KEY") or ""
+
+    headers = {"Content-Type": "application/json"}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    payload: dict = {"model": model, "messages": messages}
+    if extra:
+        payload.update(extra)
+
+    verify = os.getenv("LLM_API_VERIFY", "true").lower() in ("true", "1", "yes")
+
+    try:
+        async with httpx.AsyncClient(timeout=timeout, verify=verify) as client:
+            resp = await client.post(url, headers=headers, json=payload)
+            resp.raise_for_status()
+            data = resp.json()
+            logger.info("LLM response: %s", data)
+            return data["choices"][0]["message"]["content"]
+    except httpx.HTTPError as e:
+        logger.error("LLM API request failed (url=%s, model=%s): %s", url, model, e, exc_info=True)
+        return None
diff --git a/tests/cogs/test_ai.py b/tests/cogs/test_ai.py
index 6d92f55..aab9861 100644
--- a/tests/cogs/test_ai.py
+++ b/tests/cogs/test_ai.py
@@ -1,7 +1,5 @@
 import pytest
 from unittest.mock import AsyncMock, patch, MagicMock
-import discord
-import httpx
 from conftest import MockApplicationContext
 
 
@@ -12,36 +10,14 @@ def ai_cog(bot):
     return cog
 
 
-@pytest.fixture
-def mock_ai_response():
-    return {
-        "choices": [
-            {
-                "message": {
-                    "content": "Test response from AI"
-                }
-            }
-        ]
-    }
-
-
 @pytest.mark.asyncio
-async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot, mock_ai_response):
+async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot):
     ctx = MockApplicationContext(bot=bot, channel=mock_channel, guild=mock_guild)
 
-    mock_response = MagicMock()
-    mock_response.json = MagicMock(return_value=mock_ai_response)
-
-    mock_client = MagicMock()
-    mock_client.post = AsyncMock(return_value=mock_response)
-    mock_client.__aenter__ = AsyncMock(return_value=mock_response)
-    mock_client.__aexit__ = AsyncMock(return_value=None)
-
-    with patch('cogs.ai.httpx.AsyncClient') as mock_httpx:
-        instance = mock_httpx.return_value.__aenter__.return_value
-        instance.post = AsyncMock(return_value=mock_response)
+    with patch('cogs.ai.llm_completion', new_callable=AsyncMock, return_value="Test response from AI") as mock_llm:
         with patch.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
             await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
+            mock_llm.assert_awaited_once()
             mock_followup.assert_called_once_with("Test response from AI")
 
 
@@ -49,19 +25,10 @@ async def test_senator_command_success(ai_cog, mock_channel, mock_guild, bot, mo
 async def test_senator_command_empty_choices(ai_cog, mock_channel, mock_guild, bot):
     ctx = MockApplicationContext(bot=bot, channel=mock_channel, guild=mock_guild)
 
-    mock_response = MagicMock()
-    mock_response.json = MagicMock(return_value={"choices": []})
-
-    mock_client = MagicMock()
-    mock_client.post = AsyncMock(return_value=mock_response)
-    mock_client.__aenter__ = AsyncMock(return_value=mock_response)
-    mock_client.__aexit__ = AsyncMock(return_value=None)
-
-    with patch('cogs.ai.httpx.AsyncClient') as mock_httpx:
-        instance = mock_httpx.return_value.__aenter__.return_value
-        instance.post = AsyncMock(return_value=mock_response)
+    with patch('cogs.ai.llm_completion', new_callable=AsyncMock, return_value=None) as mock_llm:
         with patch.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
             await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
+            mock_llm.assert_awaited_once()
             mock_followup.assert_called_once_with("bot siem zesroł")
 
 
@@ -69,15 +36,8 @@ async def test_senator_command_empty_choices(ai_cog, mock_channel, mock_guild, b
 async def test_senator_command_no_choices_key(ai_cog, mock_channel, mock_guild, bot):
     ctx = MockApplicationContext(bot=bot, channel=mock_channel, guild=mock_guild)
 
-    mock_response = MagicMock()
-    mock_response.json = MagicMock(return_value={})
-
-    with patch('cogs.ai.httpx.AsyncClient') as mock_httpx:
-        instance = mock_httpx.return_value.__aenter__.return_value
-        instance.post = AsyncMock(return_value=mock_response)
-        with patch('cogs.ai.logger') as mock_logger:
-            try:
-                await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
-            except KeyError:
-                pass
-            assert mock_logger.info.called
+    with patch('cogs.ai.llm_completion', new_callable=AsyncMock, return_value=None) as mock_llm:
+        with patch.object(ctx, 'send_followup', new_callable=AsyncMock) as mock_followup:
+            await ai_cog.senator.callback(ai_cog, ctx, text="Hello", tts=False)
+            mock_llm.assert_awaited_once()
+            mock_followup.assert_called_once_with("bot siem zesroł")