RareMojo · RareMojo · Oct 5, 2023 · Oct 7, 2023 · Oct 7, 2023
diff --git a/cogs/AskDB/__init__.py b/cogs/AskDB/__init__.py
diff --git a/cogs/AskDB/askdb_cog.py b/cogs/AskDB/askdb_cog.py
@@ -1,3 +1,4 @@
+import os
 import sys
 from typing import TYPE_CHECKING
 
@@ -21,7 +22,6 @@
 handler = MongoDBHandler("askdb")
 
 
-
 class AskDB(commands.Cog):
     """
     Cog for querying documents.
@@ -34,13 +34,14 @@ def __init__(self, bot: "Bot"):
           bot (Bot): The bot instance.
         """
         self.bot = bot
+        self.default_db_id = self.bot.default_db_id
 
     @commands.hybrid_command()
     async def askdb(
         self,
         ctx: commands.Context,
         query: str,
-        db_id: str = "ba8e1813-627c-4c82-9de3-c3cfeef3d6f3",  # default to the db_id of gpt-engineer
+        db_id: str = None
     ):
         """
         Queries documents for a given query.
@@ -51,17 +52,20 @@ async def askdb(
         Returns:
         discord.Embed: An embed containing the query results.
         Examples:
-        >>> await ctx.send(embed=askdb("What is GPT-Engineer?", "2349f359-9c6e-4436-b707-af6492ddd2d7"))
-        Embed containing query results.
+        >>> /askdb What is GPT-Engineer? 2349f359-9c6e-4436-b707-af6492ddd2d7
         """
         channel = ctx.channel
         if channel.category.id != self.bot.chatbot_category_id:
             await ctx.send(embed=discord.Embed(title="Error", color=embed_color_failure, description="Please use this command in the 'AI' text-chat category."), ephemeral=True)
             return
-
+
+        if db_id == None:
+            db_id = self.default_db_id
+
         if not handler.check_exists(db_id=db_id):
             await ctx.send(embed=discord.Embed(title="Error", color=embed_color_failure, description="The DB ID you provided does not exist."), ephemeral=True)
             return
+
         await ctx.defer(ephemeral=True)
         chat_history = []
         log_debug(self.bot, f"Query: {query}")
@@ -98,6 +102,15 @@ async def askdb(
                 color=embed_color_chat,
             )
 
+            if self.bot.show_source_documents == True:
+                for i, doc in enumerate(parsed_documents, start=1):
+                    source = doc['metadata']['source']
+                    source_name = os.path.basename(source)
+                    embed.add_field(name=f"Source {i}", value=f'{source_name}', inline=True)
+
+                source_count = len(parsed_documents)
+                embed.set_footer(text=f"Total Sources: {source_count}")
+
             embed.add_field(name="Prompt:", value=f"**{query}**", inline=False)
         except Exception as e:
             log_error(self.bot, f"Error querying the DB: {e}")

diff --git a/cogs/AskDB/ingestdb_cog.py → cogs/AskDB/ingestany_cog.py b/cogs/AskDB/ingestdb_cog.py → cogs/AskDB/ingestany_cog.py
@@ -6,12 +6,10 @@
 import discord
 from discord.ext import commands
 
-from utils.ingest import ingest
+from utils.ingest import ingestAny
 from utils.mongo_db import MongoDBHandler
 from discord_bot.logger import log_debug, log_error, log_info
 
-from urllib.parse import urlparse
-
 if TYPE_CHECKING:
     from discord_bot.bot import Bot
 
@@ -23,7 +21,7 @@
 handler = MongoDBHandler("askdb")
 
 
-class IngestDBCog(commands.Cog):
+class IngestAnyCog(commands.Cog):
     """
     Cog for ingesting URLs.
     """
@@ -37,7 +35,7 @@ def __init__(self, bot: "Bot"):
         self.bot = bot
 
     @commands.hybrid_command()
-    async def ingestdb(self, ctx: commands.Context, url: str, db_name: str):
+    async def ingestany(self, ctx: commands.Context, url: str, db_name: str):
         """
         Ingests a URL.
         Args:
@@ -62,10 +60,6 @@ async def ingestdb(self, ctx: commands.Context, url: str, db_name: str):
             await ctx.send(embed=discord.Embed(title="Error", color=embed_color_failure, description="Please use this command in the 'AI' text-chat category."), ephemeral=True)
             return
 
-        parsed_url = urlparse(url)
-        if not parsed_url.netloc.endswith('readthedocs.io'):
-            await ctx.send(embed=discord.Embed(title="Error", color=embed_color_failure, description="The URL you provided is not a ReadTheDocs URL."), ephemeral=True)
-            return
         await ctx.defer(ephemeral=True)
         try:
             try:
@@ -80,7 +74,7 @@ async def ingestdb(self, ctx: commands.Context, url: str, db_name: str):
                 log_debug(
                     self.bot, f"Ingesting {url} as {db_name} for {ctx.author.name}"
                 )
-                await ingest(self.bot, url=url, namespace=random_uuid)
+                await ingestAny(self.bot, url=url, namespace=random_uuid)
                 current_time = datetime.now()
                 handler.handle_data(
                     user_id=str(ctx.author.id),
@@ -121,7 +115,7 @@ async def ingestdb(self, ctx: commands.Context, url: str, db_name: str):
 async def setup(bot: "Bot") -> None:
     """Loads the cog."""
     try:
-        await bot.add_cog(IngestDBCog(bot))
-        log_debug(bot, "IngestDBCog loaded.")
+        await bot.add_cog(IngestAnyCog(bot))
+        log_debug(bot, "IngestAnyCog loaded.")
     except Exception as e:
-        log_error(bot, f"Error loading IngestDBCog: {e}")
+        log_error(bot, f"Error loading IngestAnyCog: {e}")
diff --git a/cogs/AskDB/ingestrtd_cog.py b/cogs/AskDB/ingestrtd_cog.py
@@ -0,0 +1,121 @@
+import sys
+import uuid
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+import discord
+from discord.ext import commands
+
+from utils.ingest import ingestRTD
+from utils.mongo_db import MongoDBHandler
+from discord_bot.logger import log_debug, log_error, log_info
+
+if TYPE_CHECKING:
+    from discord_bot.bot import Bot
+
+embed_color_pending = 0xFD7C42
+embed_color_success = discord.Color.brand_green()
+embed_color_failure = discord.Color.brand_red()
+
+sys.path.append("../")
+handler = MongoDBHandler("askdb")
+
+
+class IngestRTDCog(commands.Cog):
+    """
+    Cog for ingesting URLs.
+    """
+
+    def __init__(self, bot: "Bot"):
+        """
+        Initializes the IngestDBCog class.
+        Args:
+          bot (Bot): The Bot instance.
+        """
+        self.bot = bot
+
+    @commands.hybrid_command()
+    async def ingestrtd(self, ctx: commands.Context, url: str, db_name: str):
+        """
+        Ingests a URL.
+        Args:
+        ctx (commands.Context): The context of the command.
+        url (str): The URL to ingest.
+        db_name (str): The name of the db.
+        Returns:
+        discord.Embed: An embed containing the result of the ingestion.
+        Examples:
+        >>> await ctx.send(embed=await ingestdb(ctx, 'https://example.com', 'Example db'))
+        Embed containing the result of the ingestion.
+        """
+        channel = ctx.channel
+        allowed_roles = ["Contributor", "Moderator", "Administrator", "Developer", "Head Developer", "Super Admin", "BOT"]
+        author_roles = [role.name for role in ctx.author.roles]
+
+        if not any(role in allowed_roles for role in author_roles):
+            await ctx.send(embed=discord.Embed(title="Error", color=embed_color_failure, description="You do not have permission to use this command."), ephemeral=True)
+            return
+
+        if channel.category.id != self.bot.chatbot_category_id:
+            await ctx.send(embed=discord.Embed(title="Error", color=embed_color_failure, description="Please use this command in the 'AI' text-chat category."), ephemeral=True)
+            return
+
+        await ctx.defer(ephemeral=True)
+        try:
+            try:
+                random_uuid = str(uuid.uuid4())
+                embed = discord.Embed(
+                    title="Ingesting URL",
+                    type="rich",
+                    description=f"\n *This might take a while*",
+                    color=embed_color_pending,
+                )
+                await ctx.send(embed=embed, ephemeral=True)
+                log_debug(
+                    self.bot, f"Ingesting {url} as {db_name} for {ctx.author.name}"
+                )
+                await ingestRTD(self.bot, url=url, namespace=random_uuid)
+                current_time = datetime.now()
+                handler.handle_data(
+                    user_id=str(ctx.author.id),
+                    user_name=str(ctx.author.name),
+                    db_name=db_name,
+                    db_id=random_uuid,
+                    ingest_url=url,
+                    ingested_time=current_time,
+                )
+                embed = discord.Embed(
+                    title="Success",
+                    description=f"{url}\n**DB Name:** `{db_name}`\n**DB ID:** `{random_uuid}`",
+                    timestamp=current_time,
+                    color=embed_color_success,
+                )
+                await ctx.send(embed=embed, ephemeral=True)
+            except Exception as e:
+                log_error(
+                    self.bot,
+                    f"Error ingesting {url} as {db_name} for {ctx.author.name}: {e}",
+                )
+                embed = discord.Embed(
+                    title="Error", description=f"Error: {e}", color=embed_color_failure
+                )
+            else:
+                pass
+        except Exception as e:
+            log_error(
+                self.bot,
+                f"Error ingesting {url} as {db_name} for {ctx.author.name}: {e}",
+            )
+            embed = discord.Embed(
+                title="Error", description=f"Error: {e}", color=embed_color_failure
+            )
+            await ctx.send(embed=embed, ephemeral=True)
+
+
+async def setup(bot: "Bot") -> None:
+    """Loads the cog."""
+    try:
+        await bot.add_cog(IngestRTDCog(bot))
+        log_debug(bot, "IngestRTDCog loaded.")
+    except Exception as e:
+        log_error(bot, f"Error loading IngestRTDCog: {e}")
diff --git a/cogs/Core/__init__.py b/cogs/Core/__init__.py
diff --git a/cogs/Core/chatbot_cog.py b/cogs/Core/chatbot_cog.py
@@ -120,6 +120,7 @@ async def on_message(self, ctx: discord.Message):
                                 )
 
                             chat_agent = chat_agents[self.channel_id]
+                            log_debug(self.bot, f"Prompt: {prompt}, User: {user}, Channel: {channel}")
                             messages = chat_agent.predict(prompt)
 
                             if not messages:

diff --git a/cogs/Core/core_cog.py b/cogs/Core/core_cog.py
@@ -1,5 +1,6 @@
 from discord.ext import commands
-from utils.tools import update_with_discord, welcome_to_bot
+from utils.tools import update_with_discord
+from discord_bot.terminal import welcome_to_bot
 from discord_bot.logger import log_debug, log_error, log_info
 
 from typing import TYPE_CHECKING

diff --git a/cogs/__init__.py b/cogs/__init__.py
diff --git a/configs/preprompt b/configs/preprompt
@@ -1 +1,5 @@
-The following is a friendly conversation between a human and an AI named {persona}. {persona} is talkative and provides lots of specific details from its context. {persona} will only help with the project, gpt-engineer. Gpt-engineer is a Python and AI powered software producer. You will not allow off topic conversations. If {persona} does not know the answer, truthly explain that you lack the context and suggest using the /askdocs command.
+You are a helpful AI assistant named {persona}. Very Important: If you don't know the answer, just say you don't know. DO NOT allow made up or fake answers.
+If the question is not related to Python code or gpt-engineer, politely respond that you are tuned to only answer questions that are related to Python code or gpt-engineer.
+Use as much detail as possible when responding.
+
+Now, let's think step by step and get this right:
diff --git a/discord_bot/__init__.py b/discord_bot/__init__.py
diff --git a/discord_bot/bot.py b/discord_bot/bot.py
@@ -17,9 +17,10 @@
 GUILD_ID = os.getenv("DISCORD_GUILD_ID")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 OPENAI_MODEL = os.getenv("OPENAI_MODEL")
-PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
-PINECONE_ENV = os.environ.get("PINECONE_ENV")
-PINECONE_INDEX = os.environ.get("PINECONE_INDEX")
+PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
+PINECONE_ENV = os.getenv("PINECONE_ENV")
+PINECONE_INDEX = os.getenv("PINECONE_INDEX")
+GOOGLE_SEARCH_API_KEY = os.getenv("GOOGLE_SEARCH_API_KEY")
 
 
 if TYPE_CHECKING:
@@ -46,6 +47,7 @@ def __init__(self, intents: "Intents", paths: dict, logger: "Logger"):
           Sets the bot's logger, paths, config file, avatar file, cogs directory, guild ID, owner ID, chatbot category ID, chatbot threads ID, Discord token, OpenAI API key, OpenAI model, Pinecone API key, Pinecone environment, and Pinecone index.
           Loads the config file.
           Sets the bot's display name.
+          Sets self.running to True.
         Examples:
           >>> bot = Bot(intents, paths, logger)
           Bot built.
@@ -67,11 +69,14 @@ def __init__(self, intents: "Intents", paths: dict, logger: "Logger"):
         self.pinecone_api_key = str(PINECONE_API_KEY)
         self.pinecone_env = str(PINECONE_ENV)
         self.pinecone_index = str(PINECONE_INDEX)
-
+        self.google_search_api_key = str(GOOGLE_SEARCH_API_KEY)
+
         with open(self.config_file, "r") as f:
             self.config = json.load(f)
 
+        self.default_db_id = self.config.get("default_db_id")
         self.display_name = self.config.get("bot_name")
+        self.show_source_documents = self.config.get("show_source_documents")
 
         super().__init__(command_prefix=self.config.get("prefix"), intents=intents)
         self.log.debug("Bot initialized.")
@@ -118,28 +123,34 @@ def stop_bot(self):
     async def load_cogs(self):
         """Loads all cogs in the cogs directory and its subdirectories."""
         self.log.debug("Loading cogs...")
-        total_loaded_extensions = 0
-        cog_name = None
+        num_loaded_cogs = 0
         try:
             for dirpath, dirnames, filenames in os.walk(self.cogs_dir):
                 loaded_extensions = []
+                cog_name = None
+
                 for filename in filenames:
+
                     if filename.endswith("cog.py"):
                         rel_path = os.path.relpath(dirpath, self.cogs_dir)
+
                         if rel_path == '.':
                             cog_name = f"cogs.{filename[:-3]}"
                         else:
                             cog_name = f"cogs.{rel_path.replace(os.sep, '.')}.{filename[:-3]}"
+
                         if cog_name in self.extensions:
                             continue
+
                         await self.load_extension(cog_name)
                         loaded_extensions.append(filename[:-3])
-                        total_loaded_extensions += 1
+                        num_loaded_cogs += 1
+
                 if loaded_extensions:
                     package_name = "Standalone Cogs" if dirpath == self.cogs_dir else os.path.basename(dirpath)
                     self.log.info("Loaded:")
                     self.log.info(f"Package Name: {package_name}")
                     self.log.info(f"Extensions: {', '.join(loaded_extensions)}")
         except Exception as e:
-            raise e
-        self.log.info(f"Loaded total {total_loaded_extensions} cogs.")
+            self.log.error(f"Failed to load extension: {e}")
+        self.log.info(f"Loaded total {num_loaded_cogs} cogs.")