POC: AI

commands/ai/ai_commands.py

@@ -0,0 +1,37 @@
+from discord import app_commands
+import discord
+from modules.ai.llm import get_response, list_models
+from utils.splitter import split_text
+
+def register_commands(
+    tree: discord.app_commands.CommandTree,
+    guilds: list[discord.Object],
+):
+    @tree.command(
+        name="llm",
+        description="Ask the AI a question",
+        guilds=guilds,
+    )
+    @app_commands.describe(question="The question you want to ask the AI")
+    async def llm(interaction: discord.Interaction, question: str, model: str = ""):
+        await interaction.response.defer()
+        try:
+            response = get_response(question, model)
+            split_response = split_text(response)
+            for chunk in split_response:
+                await interaction.followup.send(chunk)
+        except Exception as e:
+            await interaction.followup.send(f"An error occurred: {e}")
+
+    @tree.command(
+        name="llm_models",
+        description="List available models",
+        guilds=guilds,
+    )
+    async def llm_models(interaction: discord.Interaction):
+        await interaction.response.defer()
+        try:
+            models = list_models()
+            await interaction.followup.send("\n".join(models))
+        except Exception as e:
+            await interaction.followup.send(f"An error occurred: {e}")

main.py

@@ -8,6 +8,7 @@
 import datetime
 
 # user commands
+from commands.ai import ai_commands
 from commands.basic import basic_commands
 from commands.bbt_count import bbt_count
 from commands.config import config_commands
@@ -81,6 +82,7 @@ async def deployment_info(interaction: discord.Interaction):
 
 
 # * register commands the just the placetw server
+ai_commands.register_commands(tree, [placetw_guild])
 edit_entry_cmd.register_commands(tree, placetw_guild, client)
 restart.register_commands(tree, placetw_guild)
 watching.register_commands(tree, placetw_guild, client)

modules/ai/llm.py

@@ -0,0 +1,24 @@
+import requests
+from openai import OpenAI
+from openai.types.chat import ChatCompletion
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
+openai = OpenAI(
+    base_url=os.getenv("OPENAI_BASE_URL"),
+    api_key=os.getenv("OPENAI_API_KEY")
+)
+
+def list_models() -> list[str]:
+    return [model.id for model in openai.models.list().data]
+
+def get_response(prompt: str, model: str = "") -> str:
+    completion: ChatCompletion = openai.chat.completions.create(
+        model=model if model else os.getenv("DEFAULT_LLM_MODEL"),
+        messages=[
+            {"role": "user", "content": prompt}
+        ],
+        stream=False
+    )
+    return completion.choices[0].message.content

pyproject.toml

@@ -22,6 +22,7 @@ beautifulsoup4 = "^4.12.3"
 gitpython = "^3.1.43"
 pytest = "^8.2.1"
 
+openai = "^1.35.10"
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

utils/splitter.py

@@ -0,0 +1,17 @@
+
+
+def split_text(text: str, max_length: int = 2000) -> list[str]:
+  chunks = []
+  current_chunk = ""
+
+  for line in text.splitlines(keepends=True):  # keepends=True keeps the newline characters
+      if len(current_chunk) + len(line) <= max_length:
+          current_chunk += line
+      else:
+          chunks.append(current_chunk)
+          current_chunk = line
+
+  if current_chunk:  # Append any remaining text
+      chunks.append(current_chunk)
+
+  return chunks