import requests from bs4 import BeautifulSoup from llama_index import GPTSimpleVectorIndex from llama_index.readers.database import DatabaseReader from env import settings from logger import logger from .base import BaseToolSet, SessionGetter, ToolScope, tool class RequestsGet(BaseToolSet): @tool( name="Requests Get", description="A portal to the internet. " "Use this when you need to get specific content from a website." "Input should be a url (i.e. https://www.google.com)." "The output will be the text response of the GET request.", ) def get(self, url: str) -> str: """Run the tool.""" html = requests.get(url).text soup = BeautifulSoup(html) non_readable_tags = soup.find_all( ["script", "style", "header", "footer", "form"] ) for non_readable_tag in non_readable_tags: non_readable_tag.extract() content = soup.get_text("\n", strip=True) if len(content) > 300: content = content[:300] + "..." logger.debug( f"\nProcessed RequestsGet, Input Url: {url} " f"Output Contents: {content}" ) return content class WineDB(BaseToolSet): def __init__(self): db = DatabaseReader( scheme="postgresql", # Database Scheme host=settings["WINEDB_HOST"], # Database Host port="5432", # Database Port user="alphadom", # Database User password=settings["WINEDB_PASSWORD"], # Database Password dbname="postgres", # Database Name ) self.columns = ["nameEn", "nameKo", "description"] concat_columns = str(",'-',".join([f'"{i}"' for i in self.columns])) query = f""" SELECT Concat({concat_columns}) FROM wine """ documents = db.load_data(query=query) self.index = GPTSimpleVectorIndex(documents) @tool( name="Wine Recommendation", description="A tool to recommend wines based on a user's input. " "Inputs are necessary factors for wine recommendations, such as the user's mood today, side dishes to eat with wine, people to drink wine with, what things you want to do, the scent and taste of their favorite wine." "The output will be a list of recommended wines." "The tool is based on a database of wine reviews, which is stored in a database.", ) def recommend(self, query: str) -> str: """Run the tool.""" results = self.index.query(query) wine = "\n".join( [ f"{i}:{j}" for i, j in zip( self.columns, results.source_nodes[0].source_text.split("-") ) ] ) output = results.response + "\n\n" + wine logger.debug( f"\nProcessed WineDB, Input Query: {query} " f"Output Wine: {wine}" ) return output class ExitConversation(BaseToolSet): @tool( name="Exit Conversation", description="A tool to exit the conversation. " "Use this when you want to exit the conversation. " "The input should be a message that the conversation is over.", scope=ToolScope.SESSION, ) def exit(self, message: str, get_session: SessionGetter) -> str: """Run the tool.""" _, executor = get_session() del executor logger.debug(f"\nProcessed ExitConversation.") return message