i havent bothered to test the image & pdf funcs but we ball

README.md

@@ -0,0 +1,11 @@
+# Huginn
+
+A search engine in your discord client.
+
+## Roadmap
+- [X] Search & Open URL (credits: duccdev)
+- [X] Command exec
+- [X] Image & PDF recognitiion
+- [ ] Research
+- [ ] Image creation(?)
+- [ ] VC capabilities (difficult af to implement)

main.py

@@ -3,7 +3,7 @@ from google.genai import types
 from dotenv import load_dotenv
 from discord import app_commands
 from discord.ext import commands
-from tools import searxng, open_url
+from tools import searxng, open_url, run_command
 import os
 import io
 import discord
@@ -12,7 +12,7 @@ load_dotenv()
 
 client = genai.Client(api_key=os.getenv("GEM_API_KEY"))
 config = types.GenerateContentConfig(
-    tools=[searxng, open_url]
+    tools=[searxng, open_url, run_command]
 )
 intents = discord.Intents.default()
 intents.message_content = True

tools.py

@@ -1,5 +1,8 @@
 import aiohttp
+import subprocess
+import asyncio
 from markdownify import markdownify
+from google.genai import types
 
 SUPPORTED_TEXT_MIMETYPES = [
     "text/plain",
@@ -30,6 +33,13 @@ SUPPORTED_TEXT_MIMETYPES = [
     "application/x-yaml",
 ]
 
+SUPPORTED_IMAGE_DOCUMENT_MIMETYPES = [
+    "application/pdf",
+    "image/png",
+    "image/apng",
+    "image/jpeg"
+]
+
 
 async def searxng(query: str) -> list:
     """
@@ -79,12 +89,11 @@ async def searxng(query: str) -> list:
     return results
 
 
-async def open_url(url: str) -> dict:
+async def open_url(url: str) -> dict | types.Part:
     """
     Opens a URL and returns its full content (if it's HTML, it will be converted to clean Markdown).
     Use this when a `search` result's content is insufficient or when a user provides a direct URL to analyze.
     """
-
     async with aiohttp.ClientSession(
         headers={
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36",
@@ -107,13 +116,19 @@ async def open_url(url: str) -> dict:
             content_type = response.content_type.split(";")[0].strip()
             content_length = response.content_length or 0
 
-            if content_type not in SUPPORTED_TEXT_MIMETYPES:
+            if content_type not in SUPPORTED_TEXT_MIMETYPES + SUPPORTED_IMAGE_DOCUMENT_MIMETYPES:
                 return {
                     "content_type": content_type,
                     "content_length": content_length,
                     "content": None,
                 }
 
+            if content_type in SUPPORTED_IMAGE_DOCUMENT_MIMETYPES:
+                return types.Part.from_bytes(
+                    data=await response.read(),
+                    mime_type=content_type
+                )
+
             if "text/html" in content_type:
                 content = markdownify(await response.text())
                 if len(content) > 262144:
@@ -132,3 +147,26 @@ async def open_url(url: str) -> dict:
                 "content_length": content_length,
                 "content": content,
             }
+
+
+async def run_command(command: str) -> tuple[str, str, int]:
+    """
+    Runs a shell command on the host machine and captures its stdout, stderr and error code.
+    Args:
+     command: str
+    Returns:
+     tuple containing: stdout, stderr and error code (in that order)
+    """
+    process = await asyncio.create_subprocess_shell(
+        command,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE
+    )
+
+    stdout_data, stderr_data = await process.communicate()
+    stdout = stdout_data.decode().strip()
+    stderr = stderr_data.decode().strip()
+    return_code = process.returncode
+    if return_code is None:
+        raise TypeError
+    return stdout, stderr, return_code