Update app.py

app.py

@@ -1,114 +1,186 @@
-import streamlit as st
-import hashlib
-import os
-import aiohttp
-import asyncio
-import time
-from langsmith import traceable
-import random
-import discord
-from transformers import pipeline
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
-from pydantic import BaseModel
-from typing import List, Optional
-from tqdm import tqdm
-import re
-import os
-from teapotai import TeapotAI, TeapotAISettings
-import pandas as pd
-
-
-st.set_page_config(page_title="TeapotAI Discord Bot", page_icon=":robot_face:", layout="wide")
-
-
-
-wiki_data = pd.read_csv("https://docs.google.com/spreadsheets/d/1NNbdQWIfVHq09lMhVSN36_SkGu6XgmKTXgBWPyQcBpk/export?format=csv")
-documents = wiki_data['content'].str.split('\n\n').explode().reset_index(drop=True)
-documents
-
-# Intialize Teapot with documents
-teapot_ai = TeapotAI(documents=documents.to_list(), settings=TeapotAISettings(rag_num_results=7))
+# import streamlit as st
+# import hashlib
+# import os
+# import aiohttp
+# import asyncio
+# import time
+# from langsmith import traceable
+# import random
+# import discord
+# from transformers import pipeline
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# import numpy as np
+# from sklearn.metrics.pairwise import cosine_similarity
+# from pydantic import BaseModel
+# from typing import List, Optional
+# from tqdm import tqdm
+# import re
+# import os
+# from teapotai import TeapotAI, TeapotAISettings
+# import pandas as pd
+
+
+# st.set_page_config(page_title="TeapotAI Discord Bot", page_icon=":robot_face:", layout="wide")
+
+
+
+# wiki_data = pd.read_csv("https://docs.google.com/spreadsheets/d/1NNbdQWIfVHq09lMhVSN36_SkGu6XgmKTXgBWPyQcBpk/export?format=csv")
+# documents = wiki_data['content'].str.split('\n\n').explode().reset_index(drop=True)
+# documents
+
+# # Intialize Teapot with documents
+# teapot_ai = TeapotAI(documents=documents.to_list(), settings=TeapotAISettings(rag_num_results=7))
+
+# # I apologize, but I'm not qualified to provide personalized advice or recommendations.
+# # Latent caching is a feature of the training process that saves intermediate data and can be used to speed up the training.
+# # Yes
+# # I'm sorry, but I don't have information on the optimizers.
+# # 2
+# # The Concepts tab is made up of the following elements: For expert users: Embedding Training To get a general overview of the UI, please read the Quick Start Guide first.
+
+
+# def log_time(func):
+#     async def wrapper(*args, **kwargs):
+#         start_time = time.time()
+#         result = await func(*args, **kwargs)  # Make it awaitable
+#         end_time = time.time()
+#         print(f"{func.__name__} executed in {end_time - start_time:.4f} seconds")
+#         return result
+#     return wrapper
+
+
+# # @traceable
+# @log_time
+# async def handle_chat(user_input):
+#     def run_sync_query():
+#         return teapot_ai.query(user_input)
+
+#     output_text = await asyncio.to_thread(run_sync_query)
+
+#     return output_text
+
+
+# st.write("418 I'm a teapot")
+
+# DISCORD_TOKEN = os.environ.get("discord_key")
+
+# # Create an instance of Intents and enable the required ones
+# intents = discord.Intents.default()  # Default intents enable basic functionality
+# intents.messages = True  # Enable message-related events
+
+# # Create an instance of a client with the intents
+# client = discord.Client(intents=intents)
+
+# # Event when the bot has connected to the server
+# @client.event
+# async def on_ready():
+#     print(f'Logged in as {client.user}')
+
+# # Event when a message is received
+# @client.event
+# async def on_message(message):
+#     # Check if the message is from the bot itself to prevent a loop
+#     if message.author == client.user:
+#         return
+    
+#     # Exit the function if the bot is not mentioned
+#     if f'<@{client.user.id}>' not in message.content:
+#         return 
+        
+#     print(message.content)
+
+#     is_debug = "debug:" in message.content or "Debug:" in message.content
+#     async with message.channel.typing():
+#         cleaned_message=message.content.replace("debug:", "").replace("Debug:","").replace(f'<@{client.user.id}>',"")
+#         response = await handle_chat(cleaned_message)
+#         print(response)
+#         sent_message = await message.reply(response)
+    
+#         # Create a thread from the sent message
+#         if is_debug:
+#             thread = await sent_message.create_thread(name=f"""Debug Thread: '{cleaned_message[:80]}'""", auto_archive_duration=60)
+        
+#             # Send a message in the created thread
+#             await thread.send(debug_info)
 
-# I apologize, but I'm not qualified to provide personalized advice or recommendations.
-# Latent caching is a feature of the training process that saves intermediate data and can be used to speed up the training.
-# Yes
-# I'm sorry, but I don't have information on the optimizers.
-# 2
-# The Concepts tab is made up of the following elements: For expert users: Embedding Training To get a general overview of the UI, please read the Quick Start Guide first.
 
+# @st.cache_resource
+# def initialize():
+#     st.session_state["initialized"] = True
+#     client.run(DISCORD_TOKEN)
 
-def log_time(func):
-    async def wrapper(*args, **kwargs):
-        start_time = time.time()
-        result = await func(*args, **kwargs)  # Make it awaitable
-        end_time = time.time()
-        print(f"{func.__name__} executed in {end_time - start_time:.4f} seconds")
-        return result
-    return wrapper
+#     return 
 
+# initialize()
 
-# @traceable
-@log_time
-async def handle_chat(user_input):
-    def run_sync_query():
-        return teapot_ai.query(user_input)
 
-    output_text = await asyncio.to_thread(run_sync_query)
 
-    return output_text
 
+import streamlit as st
+import os
+import aiohttp
+import asyncio
+import discord
+import pandas as pd
+from teapotai import TeapotAI, TeapotAISettings
 
-st.write("418 I'm a teapot")
+st.set_page_config(page_title="TeapotAI Discord Bot", page_icon=":robot_face:", layout="wide")
 
 DISCORD_TOKEN = os.environ.get("discord_key")
 
-# Create an instance of Intents and enable the required ones
-intents = discord.Intents.default()  # Default intents enable basic functionality
-intents.messages = True  # Enable message-related events
-
-# Create an instance of a client with the intents
+# ========= CONFIG =========
+CONFIG = {
+    "OneTrainer": TeapotAI(
+        documents=pd.read_csv("https://docs.google.com/spreadsheets/d/1NNbdQWIfVHq09lMhVSN36_SkGu6XgmKTXgBWPyQcBpk/export?format=csv").content.str.split('\n\n').explode().reset_index(drop=True).to_list(), 
+        settings=TeapotAISettings(rag_num_results=7)),
+    "Teapot AI": TeapotAI(dataset=[])
+}
+
+# ========= DISCORD CLIENT =========
+intents = discord.Intents.default()
+intents.messages = True
 client = discord.Client(intents=intents)
 
-# Event when the bot has connected to the server
+async def handle_teapot_inference(server_name, user_input):
+    teapot_instance = CONFIG.get(server_name, CONFIG["Teapot AI"])
+    print(f"Using Teapot instance for server: {server_name}")
+    response, debug_info = await teapot_instance.infer(user_input)
+    return response, debug_info
+
+
 @client.event
 async def on_ready():
     print(f'Logged in as {client.user}')
 
-# Event when a message is received
+
 @client.event
 async def on_message(message):
-    # Check if the message is from the bot itself to prevent a loop
     if message.author == client.user:
         return
-    
-    # Exit the function if the bot is not mentioned
+
     if f'<@{client.user.id}>' not in message.content:
-        return 
-        
-    print(message.content)
+        return
 
-    is_debug = "debug:" in message.content or "Debug:" in message.content
+    server_name = message.guild.name if message.guild else "Teapot AI"
+    server_name ="OneTrainer"
     async with message.channel.typing():
-        cleaned_message=message.content.replace("debug:", "").replace("Debug:","").replace(f'<@{client.user.id}>',"")
-        response = await handle_chat(cleaned_message)
-        print(response)
+        cleaned_message = message.content.replace(f'<@{client.user.id}>', "").strip()
+        response, debug_info = await handle_teapot_inference(server_name, cleaned_message)
         sent_message = await message.reply(response)
-    
-        # Create a thread from the sent message
+
         if is_debug:
-            thread = await sent_message.create_thread(name=f"""Debug Thread: '{cleaned_message[:80]}'""", auto_archive_duration=60)
-        
-            # Send a message in the created thread
+            thread = await sent_message.create_thread(
+                name=f"Debug Thread: '{cleaned_message[:80]}'",
+                auto_archive_duration=60
+            )
             await thread.send(debug_info)
 
-
+# ========= STREAMLIT =========
 @st.cache_resource
 def initialize():
     st.session_state["initialized"] = True
     client.run(DISCORD_TOKEN)
+    return
 
-    return 
-
+st.write("418 I'm a teapot")
 initialize()