comprobot 1.1.2__tar.gz → 1.2.0__tar.gz

{comprobot-1.1.2 → comprobot-1.2.0}/Comprobot/functions.py

@@ -12,6 +12,8 @@ from google.genai.errors import ClientError
 from .bot import client
 from .data import ai, system_prompt_text
 
+context_memory: Dict[int, List[Dict[str, Any]]] = {}
+
 
 def para(count=1):
     for i in range(count):
@@ -61,58 +63,102 @@ def demoji(text):
     return emoji_pattern.sub("", text)
 
 
-def chat(message):
-    messages: List[Dict[str, Any]] = []
-    user_id = client.user.id if client.user else ""
+async def get_referenced_message(message):
+    if not message.reference:
+        return None
+
+    if isinstance(message.reference.resolved, discord.Message):
+        return message.reference.resolved
+
+    try:
+        return await message.channel.fetch_message(message.reference.message_id)
+    except discord.NotFound:
+        return None
+
 
-    if ai["provider"].lower() in ("ollama", "groq"):
-        messages.append(
-            {
-                "role": "user",
-                "content": message.content.replace(f"<@{user_id}>", ""),
-            }
+async def chat(message):
+    user_id = client.user.id if client.user else ""
+    channel_id = message.channel.id
+
+    if channel_id not in context_memory:
+        context_memory[channel_id] = []
+
+    messages = context_memory[channel_id]
+
+    referenced_message = await get_referenced_message(message)
+
+    # Construct user content based on templates
+    attachment_url = ""
+    if message.attachments and ai["include_attachment"]:
+        attachment_url = message.attachments[0].url
+    elif referenced_message and ai["include_attachment"]:
+        if referenced_message.attachments:
+            attachment_url = referenced_message.attachments[0].url
+
+    user_prompt_content = (
+        ai["user_prompt_structure"]
+        .replace("{{PROMPT}}", ai["user_prompt"])
+        .replace(
+            "{{REPLY_PROMPT}}",
+            ai["user_reply_prompt"]
+            if ai["include_reply"] and referenced_message
+            else "",
         )
-    elif ai["provider"].lower() == "gemini":
-        messages.append(
-            {
-                "role": "user",
-                "parts": [{"text": message.content.replace(f"<@{user_id}>", "")}],
-            }
+        .replace(
+            "{{ATTACHMENT_PROMPT}}",
+            ai["user_attachement_prompt"].replace("{{FILE}}", attachment_url)
+            if attachment_url
+            else "",
         )
-    else:
-        raise ValueError(f"Unknown provider: {ai['provider']}")
+        .replace(
+            "{{USERNAME}}",
+            message.author.display_name if ai["include_username"] else "",
+        )
+        .replace("{{MESSAGE}}", message.content.replace(f"<@{user_id}>", ""))
+        .replace(
+            "{{REPLY}}",
+            referenced_message.content
+            if referenced_message and ai["include_reply"]
+            else "",
+        )
+    )
 
-    if ai["provider"].lower() == "ollama":
-        messages.insert(
-            0,
-            {
-                "role": "system",
-                "content": system_prompt_text,
-            },
+    # Add user message to history
+    messages.append({"role": "user", "content": user_prompt_content})
+
+    # Limit context
+    max_total = cast(int, ai["max_messages_context"])
+    while len(messages) > max_total:
+        messages.pop(0)
+
+    provider = ai["provider"].lower()
+    content = ""
+
+    if provider == "ollama":
+        formatted_messages = [
+            {"role": "system", "content": system_prompt_text}
+        ] + messages
+        response = ollama.chat(
+            model=cast(str, ai["model"]), messages=formatted_messages
         )
-        response = ollama.chat(model=cast(str, ai["model"]), messages=messages)
-        content: str = response.message.content or ""
-    elif ai["provider"].lower() == "gemini":
+        content = response.message.content or ""
+    elif provider == "gemini":
         gemini_client = genai.Client(api_key=os.getenv("GEMINI"))
-        system_instruction = None
         formatted_messages = []
         for m in messages:
-            if m["role"] == "system":
-                system_instruction = m["content"]
-            else:
-                formatted_messages.append(
-                    {
-                        "role": m["role"],
-                        "parts": [
-                            {"text": m["content"] if "content" in m else m["parts"][0]}
-                        ],
-                    }
-                )
+            formatted_messages.append(
+                {
+                    "role": "user" if m["role"] == "user" else "model",
+                    "parts": [{"text": m["content"]}],
+                }
+            )
+
         config = None
-        if system_instruction:
+        if system_prompt_text:
             config = genai.types.GenerateContentConfig(
-                system_instruction=system_instruction
+                system_instruction=system_prompt_text
             )
+
         try:
             response = gemini_client.models.generate_content(
                 model=ai["model"], contents=formatted_messages, config=config
@@ -123,38 +169,25 @@ def chat(message):
                 available = gemini_client.models.list()
                 print(f"Available models: {[m.name for m in available]}")
             raise
-    elif ai["provider"].lower() == "groq":
+    elif provider == "groq":
         groq_client = groq.Groq(api_key=os.getenv("GROQ"))
         formatted_messages = [
-            {"role": "system", "content": system_prompt_text},
-            messages[0],
-        ]
+            {"role": "system", "content": system_prompt_text}
+        ] + messages
         response = groq_client.chat.completions.create(
-            model=ai["model"], messages=formatted_messages
+            model=ai["model"],
+            messages=formatted_messages,  # type: ignore
         )
         content = response.choices[0].message.content or ""
     else:
         raise ValueError(f"Unknown provider: {ai['provider']}")
 
-    if ai["provider"].lower() == "groq":
-        role = "assistant"
-    elif ai["provider"].lower() == "ollama":
-        role = "assistant"
-    elif ai["provider"].lower() == "gemini":
-        role = "model"
-    else:
-        raise ValueError(f"Unknown provider: {ai['provider']}")
-
-    messages.append(
-        {
-            "role": role,
-            "content": content,
-        }
-    )
+    # Add assistant message to history
+    messages.append({"role": "assistant", "content": content})
 
-    max_total = 1 + cast(int, ai["max_messages_context"])
+    # Prune again after adding assistant message
     while len(messages) > max_total:
-        messages.pop(1)
+        messages.pop(0)
 
     if ai["remove_emojis"]:
         content = demoji(content)

{comprobot-1.1.2 → comprobot-1.2.0}/Comprobot/main.py

@@ -66,11 +66,11 @@ async def on_message(message):
 
     if (
         f"<@{user_id}>" in message.content
-        or (is_reply_to_bot if config["answer_to_reply"] else False)
+        or (is_reply_to_bot if ai["answer_to_reply"] else False)
         and ai["activate_ai"]
     ):
         async with message.channel.typing():
-            response = chat(message)
+            response = await chat(message)
 
     if response:
         async with message.channel.typing():

{comprobot-1.1.2 → comprobot-1.2.0}/Comprobot/templates.py

@@ -6,11 +6,35 @@ max_messages_context = 10
 remove_emojis        = true
 lower_response       = true
 answer_to_reply      = false
+include_username     = true
+include_reply        = true
+include_attachment   = true
 
 system_prompt = """
 You are a helpful assistant that gives short, helpful answers.
 Your answers can maximally be 1000 characters long.
-"""'''
+"""
+
+user_prompt = """
+{{USERNAME}} wrote the following message: {{MESSAGE}}
+"""
+
+user_reply_prompt = """
+They added the following context: {{REPLY}}
+"""
+
+user_attachement_prompt = """
+The user attached the following file: {{FILE}}
+"""
+
+user_prompt_structure = """
+{{PROMPT}}
+
+{{REPLY_PROMPT}}
+
+{{ATTACHMENT_PROMPT}}
+"""
+'''
 
 config = r"""
 command_prefix =  "!"

{comprobot-1.1.2/comprobot.egg-info → comprobot-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comprobot
-Version: 1.1.2
+Version: 1.2.0
 Summary: A self-hostable Discord bot built for maximum customization.
 Author: badluma
 License: MIT

{comprobot-1.1.2 → comprobot-1.2.0/comprobot.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: comprobot
-Version: 1.1.2
+Version: 1.2.0
 Summary: A self-hostable Discord bot built for maximum customization.
 Author: badluma
 License: MIT

{comprobot-1.1.2 → comprobot-1.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "comprobot"
-version = "1.1.2"
+version = "1.2.0"
 authors = [{name = "badluma"}]
 description = "A self-hostable Discord bot built for maximum customization."
 readme = "README.md"