Update app.py

app.py

@@ -12,25 +12,28 @@ st.set_page_config(page_title="TeapotAI Discord Bot", page_icon=":robot_face:",
 
 DISCORD_TOKEN = os.environ.get("discord_key")
 
-# ========= CONFIG =========
-CONFIG = {
-    # "OneTrainer": TeapotAI(
-    #     documents=pd.read_csv("https://docs.google.com/spreadsheets/d/1NNbdQWIfVHq09lMhVSN36_SkGu6XgmKTXgBWPyQcBpk/export?gid=361556791&format=csv").content.str.split('\n\n').explode().reset_index(drop=True).to_list(), 
-    #     settings=TeapotAISettings(rag_num_results=7)
-    # ),
-    "Teapot AI": TeapotAI(
-        documents=pd.read_csv("https://docs.google.com/spreadsheets/d/1NNbdQWIfVHq09lMhVSN36_SkGu6XgmKTXgBWPyQcBpk/export?gid=1617599323&format=csv").content.str.split('\n\n').explode().reset_index(drop=True).to_list(), 
-        settings=TeapotAISettings(rag_num_results=3)
-    ),
-}
 
-#  ========= SEARCH API  =========
+#  ======= API KEYS =======
 
-API_KEY = os.environ.get("brave_api_key")
+BRAVE_API_KEY = os.environ.get("brave_api_key")
+WEATHER_API_KEY = os.environ.get("weather_api_key")
+
+# ======== TOOLS ===========
+import requests
+from typing import Optional
+from teapotai import TeapotTool
+import re
+import math
+import pandas as pd
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, logging
+
+### SEARCH TOOL
+class BraveWebSearch(BaseModel):
+    search_query: str = Field(..., description="the search string to answer the question")
 
 def brave_search_context(query, count=3):
     url = "https://api.search.brave.com/res/v1/web/search"
-    headers = {"Accept": "application/json", "X-Subscription-Token": API_KEY}
+    headers = {"Accept": "application/json", "X-Subscription-Token": BRAVE_API_KEY}
     params = {"q": query, "count": count}
     
     response = requests.get(url, headers=headers, params=params)
@@ -43,6 +46,155 @@ def brave_search_context(query, count=3):
         print(f"Error: {response.status_code}, {response.text}")
         return ""
 
+### CALCULATOR TOOL
+def evaluate_expression(expr) -> str:
+    """
+    Evaluate a simple algebraic expression string safely.
+    Supports +, -, *, /, **, and parentheses.
+    Retries evaluation after stripping non-numeric/non-operator characters if needed.
+    """
+
+    expr = expr.expression
+
+    allowed_names = {
+        k: v for k, v in vars(__builtins__).items()
+        if k in ("abs", "round")
+    }
+
+    allowed_names.update({k: getattr(math, k) for k in ("sqrt", "pow")})
+
+    def safe_eval(expression):
+        return eval(expression, {"__builtins__": None}, allowed_names)
+
+    try:
+        result = safe_eval(expr)
+        return f"{expr} = {result}"
+    except Exception as e:
+        print(f"Initial evaluation failed: {e}")
+        # Strip out any characters that are not numbers, parentheses, or valid operators
+        cleaned_expr = re.sub(r"[^0-9\.\+\-\*/\*\*\(\) ]", "", expr)
+        try:
+            result = safe_eval(cleaned_expr)
+            return f"{cleaned_expr} = {result}"
+        except Exception as e2:
+            print(f"Retry also failed: {e2}")
+            return "Sorry, I am unable to calculate that."
+
+
+class Calculator(BaseModel):
+    expression: str = Field(..., description="mathematical expression")
+
+
+### Weather Tool
+def get_weather(city_name):
+    # OpenWeatherMap API endpoint
+    url = f'https://api.openweathermap.org/data/2.5/weather?appid={WEATHER_API_KEY}&units=imperial&q={city_name}'
+    
+
+    # Send GET request to the OpenWeatherMap API
+    response = requests.get(url)
+    
+    # Check if the request was successful
+    if response.status_code == 200:
+        data = response.json()
+        
+        # Extract relevant weather information
+        city = data['name']
+        temperature = round(data['main']['temp'])
+        weather_description = data['weather'][0]['description']
+        
+        # Print or return the results
+        return f"The weather in {city} is {weather_description} with a temperature of {temperature}°F."
+    else:
+        return "City not found or there was an error with the request."
+
+class Weather(BaseModel):
+    city_name: str = Field(..., description="The name of the city to pull the weather for")
+    
+def get_weather(weather_schema):
+    # OpenWeatherMap API endpoint
+    url = f'https://api.openweathermap.org/data/2.5/weather?appid=de016a1d30e7bbe278971c2b17aabca0&units=imperial&q={weather_schema.city_name}'
+    
+
+    # Send GET request to the OpenWeatherMap API
+    response = requests.get(url)
+    
+    # Check if the request was successful
+    if response.status_code == 200:
+        data = response.json()
+        
+        # Extract relevant weather information
+        city = data['name']
+        temperature = round(data['main']['temp'])
+        weather_description = data['weather'][0]['description']
+        
+        # Print or return the results
+        return f"The weather in {city} is {weather_description} with a temperature of {temperature}°F."
+    else:
+        return "City not found or there was an error with the request."
+
+
+### Stupid Question Tool
+class CountNumberLetter(BaseModel):
+    word: str = Field(..., description="the word to count the number of letters in")
+    letter: str = Field(..., description="the letter to count the occurences of")
+
+def count_number_letters(obj):
+  letter = obj.letter.lower()
+  expression = obj.word.lower()
+  count = len([l for l in expression if l == letter])
+  if count == 1:
+    return f"There is 1 '{letter}' in '{expression}'"
+  return f"There are {count} '{letter}'s in '{expression}'"
+
+
+
+DEFAULT_TOOLS = [
+    TeapotTool(
+        name="websearch",
+        description="Execute web searches with pagination and filtering",
+        schema=BraveWebSearch,
+        fn=brave_search_context
+    ),
+    TeapotTool(
+        name="calculator",
+        description="Perform calculations",
+        schema=Calculator,
+        fn=lambda expression: evaluate_expression(expression),
+    ),
+    TeapotTool(
+        name="letter_counter",
+        description="Can count how many times a letter occurs in a word.",
+        schema=CountNumberLetter,
+        fn=count_number_letters
+    ),
+    TeapotTool(
+        name="weather",
+        description="Can pull today's weather information for any city.",
+        schema=Weather,
+        fn=get_weather
+    )
+]
+
+# ========= CONFIG =========
+CONFIG = {
+    # "OneTrainer": TeapotAI(
+    #     documents=pd.read_csv("https://docs.google.com/spreadsheets/d/1NNbdQWIfVHq09lMhVSN36_SkGu6XgmKTXgBWPyQcBpk/export?gid=361556791&format=csv").content.str.split('\n\n').explode().reset_index(drop=True).to_list(), 
+    #     settings=TeapotAISettings(rag_num_results=7)
+    # ),
+    "Teapot AI": TeapotAI(
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+          "teapotai/teapotllm",
+          revision="5aa6f84b5bd59da85552d55cc00efb702869cbf8",
+        ),
+        documents=pd.read_csv("https://docs.google.com/spreadsheets/d/1NNbdQWIfVHq09lMhVSN36_SkGu6XgmKTXgBWPyQcBpk/export?gid=1617599323&format=csv").content.str.split('\n\n').explode().reset_index(drop=True).to_list(), 
+        settings=TeapotAISettings(rag_num_results=3, log_level="debug"),
+        tools=DEFAULT_TOOLS
+    ),
+}
+
+
+
 # ========= DISCORD CLIENT =========
 intents = discord.Intents.default()
 intents.messages = True