mirror of https://github.com/corca-ai/EVAL
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
108 lines
3.5 KiB
Python
108 lines
3.5 KiB
Python
1 year ago
|
from env import settings
|
||
|
|
||
|
import requests
|
||
|
|
||
|
from llama_index.readers.database import DatabaseReader
|
||
|
from llama_index import GPTSimpleVectorIndex
|
||
|
|
||
1 year ago
|
from bs4 import BeautifulSoup
|
||
1 year ago
|
|
||
1 year ago
|
from .base import tool, BaseToolSet, ToolScope, SessionGetter
|
||
1 year ago
|
from logger import logger
|
||
1 year ago
|
|
||
1 year ago
|
|
||
|
class RequestsGet(BaseToolSet):
|
||
|
@tool(
|
||
1 year ago
|
name="Requests Get",
|
||
1 year ago
|
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.",
|
||
|
)
|
||
1 year ago
|
def get(self, url: str) -> str:
|
||
1 year ago
|
"""Run the tool."""
|
||
1 year ago
|
html = requests.get(url).text
|
||
|
soup = BeautifulSoup(html)
|
||
|
non_readable_tags = soup.find_all(
|
||
|
["script", "style", "header", "footer", "form"]
|
||
|
)
|
||
1 year ago
|
|
||
1 year ago
|
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] + "..."
|
||
1 year ago
|
|
||
1 year ago
|
logger.debug(
|
||
1 year ago
|
f"\nProcessed RequestsGet, Input Url: {url} " f"Output Contents: {content}"
|
||
|
)
|
||
|
|
||
1 year ago
|
return content
|
||
1 year ago
|
|
||
|
|
||
1 year ago
|
class WineDB(BaseToolSet):
|
||
1 year ago
|
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)
|
||
|
|
||
1 year ago
|
@tool(
|
||
1 year ago
|
name="Wine Recommendation",
|
||
1 year ago
|
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.",
|
||
|
)
|
||
1 year ago
|
def recommend(self, query: str) -> str:
|
||
1 year ago
|
"""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("-")
|
||
|
)
|
||
|
]
|
||
|
)
|
||
1 year ago
|
output = results.response + "\n\n" + wine
|
||
|
|
||
1 year ago
|
logger.debug(
|
||
|
f"\nProcessed WineDB, Input Query: {query} " f"Output Wine: {wine}"
|
||
|
)
|
||
1 year ago
|
|
||
|
return output
|
||
1 year ago
|
|
||
|
|
||
1 year ago
|
class ExitConversation(BaseToolSet):
|
||
|
@tool(
|
||
1 year ago
|
name="Exit Conversation",
|
||
1 year ago
|
description="A tool to exit the conversation. "
|
||
1 year ago
|
"Use this when you want to exit the conversation. "
|
||
1 year ago
|
"The input should be a message that the conversation is over.",
|
||
1 year ago
|
scope=ToolScope.SESSION,
|
||
1 year ago
|
)
|
||
1 year ago
|
def exit(self, message: str, get_session: SessionGetter) -> str:
|
||
1 year ago
|
"""Run the tool."""
|
||
1 year ago
|
_, executor = get_session()
|
||
|
del executor
|
||
1 year ago
|
|
||
1 year ago
|
logger.debug(f"\nProcessed ExitConversation.")
|
||
1 year ago
|
|
||
1 year ago
|
return message
|