parent
83044a905b
commit
ba7eaaed2a
19 changed files with 3189 additions and 584 deletions
@ -1,314 +0,0 @@ |
||||
|
||||
import streamlit as st |
||||
from time import sleep |
||||
from datetime import datetime, timedelta |
||||
from colorprinter.print_color import * |
||||
|
||||
|
||||
from _base_class import StreamlitBaseClass |
||||
from _rss import RSSReader |
||||
from projects_page import Project |
||||
from streamlit_chatbot import StreamlitChat |
||||
|
||||
|
||||
class BotChatPage(StreamlitBaseClass): |
||||
def __init__(self, username): |
||||
super().__init__(username=username) |
||||
self.collection_name = None |
||||
self.project_name = None |
||||
self.project: Project = None |
||||
self.chat = None |
||||
self.role = "Research Assistant" # Default persona |
||||
self.page_name = "Bot Chat" |
||||
self.chat_key = None |
||||
|
||||
# Initialize attributes from session state if available |
||||
if self.page_name in st.session_state: |
||||
for k, v in st.session_state[self.page_name].items(): |
||||
setattr(self, k, v) |
||||
|
||||
def run(self): |
||||
from streamlit_chatbot import EditorBot, ResearchAssistantBot, PodBot, StreamlitBot |
||||
bot = None |
||||
self.update_current_page("Bot Chat") |
||||
self.remove_old_unsaved_chats() |
||||
self.sidebar_actions() |
||||
|
||||
if self.collection_name or self.project: |
||||
print_purple("Collection:", self.collection_name, "Project:", self.project_name) |
||||
# If no chat exists, create a new Chat instance |
||||
self.chat = self.get_chat(role=self.role) |
||||
|
||||
# Create a Bot instance with the Chat object |
||||
if self.role == "Research Assistant": |
||||
print_blue("Creating Research Assistant Bot") |
||||
bot = ResearchAssistantBot( |
||||
username=self.username, |
||||
chat=self.chat, |
||||
collection=self.collection_name, |
||||
project=self.project, |
||||
) |
||||
|
||||
elif self.role == "Editor": |
||||
bot = EditorBot( |
||||
username=self.username, |
||||
chat=self.chat, |
||||
collection=self.collection, |
||||
project=self.project, |
||||
) |
||||
elif self.role == "Podcast": |
||||
st.session_state["make_podcast"] = True |
||||
# with st.sidebar: |
||||
with st.sidebar: |
||||
with st.form("make_podcast_form"): |
||||
instructions = st.text_area( |
||||
"What should the podcast be about? Give a brief description, as if you were the producer." |
||||
) |
||||
start = st.form_submit_button("Make Podcast!") |
||||
if start: |
||||
bot = PodBot( |
||||
subject=self.project.name, |
||||
username=self.username, |
||||
chat=self.chat, |
||||
collection=self.collection, |
||||
project=self.project, |
||||
instructions=instructions, |
||||
) |
||||
|
||||
# Run the bot (this will display chat history and process user input) |
||||
if bot: |
||||
bot.run() |
||||
|
||||
# Save updated chat state to session state |
||||
st.session_state[self.page_name] = { |
||||
"collection": self.collection, |
||||
"project": self.project, |
||||
"chat": self.chat, |
||||
"role": self.role, |
||||
} |
||||
else: # If no collection or project is selected, use the conversational response bot |
||||
print_yellow("No collection or project selected. Using conversational response bot.") |
||||
bot = StreamlitBot( |
||||
username=self.username, |
||||
chat=self.get_chat(), |
||||
tools=["conversational_response_tool"], |
||||
) |
||||
bot.run() |
||||
|
||||
|
||||
def get_chat(self, role="Research Assistant"): |
||||
print_blue('CHAT TYPE:', role) |
||||
if 'chat_key' not in st.session_state: |
||||
chat = StreamlitChat(username=self.username, role=role) |
||||
st.session_state['chat_key'] = chat._key |
||||
print_blue("Creating new chat:", st.session_state['chat_key']) |
||||
else: |
||||
print_blue("Old chat:", st.session_state['chat_key']) |
||||
chat_data = self.user_arango.db.collection("chats").get(st.session_state['chat_key']) |
||||
chat = StreamlitChat.from_dict(chat_data) |
||||
return chat |
||||
|
||||
def sidebar_actions(self): |
||||
with st.sidebar: |
||||
self.collection = self.choose_collection( |
||||
"Article collection to use for chat:" |
||||
) |
||||
self.project = self.choose_project("Project to use for chat:") |
||||
|
||||
if self.collection or self.project: |
||||
st.write("---") |
||||
if self.project: |
||||
self.role = st.selectbox( |
||||
"Choose Bot Role", |
||||
options=["Research Assistant", "Editor", "Podcast"], |
||||
index=0, |
||||
) |
||||
elif self.collection: |
||||
self.role = "Research Assistant" |
||||
|
||||
# Load existing chats from the database |
||||
if self.project: |
||||
chat_history = list( |
||||
self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats FILTER doc["project"] == "{self.project}" RETURN {{"_key": doc["_key"], "name": doc["name"]}}' |
||||
) |
||||
) |
||||
# self.project = Project(username=self.username, project_name=self.project_name, user_arango=self.user_arango) |
||||
elif self.collection: |
||||
chat_history = list( |
||||
self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats FILTER doc["collection"] == "{self.collection}" RETURN {{"_key": doc["_key"], "name": doc["name"]}}' |
||||
) |
||||
) |
||||
|
||||
chats = {i["name"]: i["_key"] for i in chat_history} |
||||
selected_chat = st.selectbox( |
||||
"Continue another chat", options=[""] + list(chats.keys()), index=0 |
||||
) |
||||
if selected_chat: |
||||
st.session_state["chat_key"] = chats[selected_chat] |
||||
self.chat = self.get_chat() |
||||
|
||||
if not self.role: |
||||
self.role == "Research Assistant" |
||||
|
||||
def remove_old_unsaved_chats(self): |
||||
two_weeks_ago = datetime.now() - timedelta(weeks=2) |
||||
q = f'FOR doc IN chats FILTER doc.saved == false AND doc.last_updated < "{two_weeks_ago.isoformat()}" RETURN doc' |
||||
print_blue(q) |
||||
old_chats = self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats RETURN doc' |
||||
) |
||||
print('test', old_chats) |
||||
old_chats = self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats FILTER doc.saved == false AND doc.last_updated < "{two_weeks_ago.isoformat()}" RETURN doc' |
||||
) |
||||
for chat in old_chats: |
||||
print_red(chat["_id"]) |
||||
self.user_arango.db.collection("chats").delete(chat["_key"]) |
||||
|
||||
|
||||
class SettingsPage(StreamlitBaseClass): |
||||
def __init__(self, username: str): |
||||
super().__init__(username=username) |
||||
|
||||
def run(self): |
||||
self.update_current_page("Settings") |
||||
self.set_profile_picture() |
||||
self.use_reasoning_model() |
||||
|
||||
def set_profile_picture(self): |
||||
st.markdown("Profile picture") |
||||
profile_picture = st.file_uploader( |
||||
"Upload profile picture", type=["png", "jpg", "jpeg"] |
||||
) |
||||
if profile_picture: |
||||
# Resize the image to 64x64 pixels |
||||
from PIL import Image |
||||
|
||||
img = Image.open(profile_picture) |
||||
img.thumbnail((64, 64)) |
||||
img_path = f"user_data/{st.session_state['username']}/profile_picture.png" |
||||
img.save(img_path) |
||||
self.update_settings("avatar", img_path) |
||||
st.success("Profile picture uploaded") |
||||
sleep(1) |
||||
|
||||
def use_reasoning_model(self): |
||||
""" |
||||
Displays a checkbox in the Streamlit interface to enable or disable the reasoning model for generating responses in chats. |
||||
|
||||
Retrieves the current settings and checks if the "use_reasoning_model" key exists. If not, it initializes it to False. |
||||
Then, it displays a markdown text and a checkbox for the user to toggle the reasoning model usage. |
||||
The updated setting is saved back to the settings. |
||||
|
||||
Returns: |
||||
None |
||||
""" |
||||
settings = self.get_settings() |
||||
if "use_reasoning_model" not in settings: |
||||
settings["use_reasoning_model"] = False |
||||
st.markdown("Use Reasoning Model") |
||||
|
||||
use_reasoning_model = st.checkbox("Use Reasoning Model", value=settings["use_reasoning_model"], help="Use the reasoning model to generate responses in chats. This may take longer to process.") |
||||
self.update_settings("use_reasoning_model", use_reasoning_model) |
||||
|
||||
|
||||
class RSSFeedsPage(StreamlitBaseClass): |
||||
def __init__(self, username: str): |
||||
super().__init__(username=username) |
||||
self.page_name = "RSS Feeds" |
||||
self.reader = RSSReader(username=username) |
||||
# Initialize attributes from session state if available |
||||
for k, v in st.session_state.get(self.page_name, {}).items(): |
||||
setattr(self, k, v) |
||||
|
||||
def run(self): |
||||
if "selected_feed" not in st.session_state: |
||||
st.session_state["selected_feed"] = None |
||||
self.update_current_page(self.page_name) |
||||
self.display_feed() |
||||
self.sidebar_actions() |
||||
self.update_session_state(page_name=self.page_name) |
||||
|
||||
def select_rss_feeds(self): |
||||
rss_feeds = self.reader.get_rss_feeds() |
||||
if rss_feeds: |
||||
feed_options = [feed["title"] for feed in rss_feeds] |
||||
with st.sidebar: |
||||
st.subheader("Show your feeds") |
||||
selected_feed_title = st.selectbox( |
||||
"Select a feed", options=feed_options, index=None |
||||
) |
||||
if selected_feed_title: |
||||
st.session_state["selected_feed"] = [ |
||||
feed["_key"] |
||||
for feed in rss_feeds |
||||
if feed["title"] == selected_feed_title |
||||
][0] |
||||
st.rerun() |
||||
else: |
||||
st.write("You have no RSS feeds added.") |
||||
|
||||
def search_feeds(self, rss_url): |
||||
with st.spinner("Discovering feeds..."): |
||||
feeds = self.reader.discover_feeds(rss_url) |
||||
if feeds: |
||||
st.session_state["discovered_feeds"] = feeds |
||||
else: |
||||
st.error("No RSS feeds found at the provided URL.") |
||||
|
||||
def sidebar_actions(self): |
||||
if "discovered_feeds" not in st.session_state: |
||||
st.session_state["discovered_feeds"] = None |
||||
|
||||
with st.sidebar: |
||||
self.select_rss_feeds() |
||||
st.subheader("Add a New RSS Feed") |
||||
with st.form("add_rss_feed"): |
||||
rss_url = st.text_input("Website URL or RSS Feed URL") |
||||
submitted = st.form_submit_button("Discover Feeds") |
||||
if submitted: |
||||
print_green(rss_url) |
||||
feeds = self.reader.discover_feeds(rss_url) |
||||
st.session_state["discovered_feeds"] = feeds |
||||
|
||||
if st.session_state["discovered_feeds"]: |
||||
st.subheader("Select a Feed to Add") |
||||
feeds = st.session_state["discovered_feeds"] |
||||
feed_options = [f"{feed['title']} ({feed['href']})" for feed in feeds] |
||||
selected_feed = st.selectbox("Available Feeds", options=feed_options) |
||||
selected_feed_url = feeds[feed_options.index(selected_feed)]["href"] |
||||
|
||||
if st.button("Preview Feed"): |
||||
feed = self.reader.parse_feed(selected_feed_url) |
||||
st.write(f"{feed.title}") |
||||
description = self.reader.html_to_markdown(feed.description) |
||||
st.write(f"_{description}_") |
||||
for entry in feed.entries[:5]: |
||||
with st.expander(entry["title"]): |
||||
summary = entry.get("summary", "No summary available") |
||||
markdown_summary = self.reader.html_to_markdown(summary) |
||||
st.markdown(markdown_summary) |
||||
print_yellow(selected_feed_url) |
||||
|
||||
if st.button( |
||||
"Add RSS Feed", |
||||
on_click=self.reader.add_rss_feed, |
||||
args=[selected_feed_url], |
||||
): |
||||
del st.session_state["discovered_feeds"] |
||||
st.success("RSS Feed added.") |
||||
st.rerun() |
||||
|
||||
def display_feed(self): |
||||
if st.session_state["selected_feed"]: |
||||
self.reader.get_feed(st.session_state["selected_feed"]) |
||||
st.title(self.reader.feed.title) |
||||
st.write(f"_{self.reader.feed.description}_") |
||||
for entry in self.reader.feed.entries[:5]: |
||||
with st.expander(entry["title"]): |
||||
summary = entry.get("summary", "No summary available") |
||||
markdown_summary = self.reader.html_to_markdown(summary) |
||||
st.markdown(markdown_summary) |
||||
st.markdown(f"[Read more]({entry['link']})") |
||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,236 @@ |
||||
|
||||
import streamlit as st |
||||
from datetime import datetime, timedelta |
||||
from colorprinter.print_color import * |
||||
|
||||
|
||||
from _base_class import StreamlitBaseClass |
||||
from _rss import RSSReader |
||||
from projects_page import Project |
||||
from streamlit_chatbot import StreamlitChat, StreamlitBot |
||||
|
||||
|
||||
class BotChatPage(StreamlitBaseClass): |
||||
""" |
||||
BotChatPage - A Streamlit interface for chatting with various AI assistants. |
||||
This class provides a user interface for interacting with different types of AI bots |
||||
(Research Assistant, Editor, Podcast) that can access and work with user's collections |
||||
and projects. |
||||
Attributes: |
||||
username (str): The username of the current user. |
||||
collection_name (str): Name of the selected collection. |
||||
project_name (str): Name of the selected project. |
||||
project (Project): Project instance the chat is associated with. |
||||
chat (StreamlitChat): Chat instance for maintaining conversation history. |
||||
role (str): The selected bot persona, default is "Research Assistant". |
||||
page_name (str): Name of the current page ("Bot Chat"). |
||||
chat_key (str): Unique identifier for the current chat session. |
||||
bot (StreamlitBot): Instance of the selected bot type. |
||||
Methods: |
||||
run(): Main method to render the chat interface and handle interactions. |
||||
get_chat(role, new_chat): Retrieves existing chat or creates a new one. |
||||
sidebar_actions(): Renders sidebar elements for selecting collections, projects, and chat options. |
||||
remove_old_unsaved_chats(): Cleans up unsaved chats older than two weeks. |
||||
""" |
||||
def __init__(self, username): |
||||
super().__init__(username=username) |
||||
self.collection_name = None |
||||
self.project_name = None |
||||
self.project: Project = None |
||||
self.chat = None |
||||
self.role = "Research Assistant" # Default persona |
||||
self.page_name = "Bot Chat" |
||||
self.chat_key = None |
||||
self.bot: StreamlitBot = None |
||||
|
||||
# Initialize attributes from session state if available |
||||
if self.page_name in st.session_state: |
||||
for k, v in st.session_state[self.page_name].items(): |
||||
setattr(self, k, v) |
||||
|
||||
def run(self): |
||||
from streamlit_chatbot import EditorBot, ResearchAssistantBot, PodBot, StreamlitBot |
||||
self.bot: StreamlitBot = None |
||||
self.update_current_page("Bot Chat") |
||||
self.remove_old_unsaved_chats() |
||||
self.sidebar_actions() |
||||
|
||||
if self.collection_name or self.project: |
||||
print_purple("Collection:", self.collection_name, "Project:", self.project_name) |
||||
# If no chat exists, create a new Chat instance |
||||
self.chat = self.get_chat(role=self.role) |
||||
|
||||
# Create a Bot instance with the Chat object |
||||
if self.role == "Research Assistant": |
||||
print_blue("Creating Research Assistant Bot") |
||||
self.bot: ResearchAssistantBot = ResearchAssistantBot( |
||||
username=self.username, |
||||
chat=self.chat, |
||||
collection=self.collection_name, |
||||
project=self.project, |
||||
tools=[ |
||||
"fetch_other_documents_tool", |
||||
"fetch_science_articles_tool", |
||||
"fetch_science_articles_and_other_documents_tool", |
||||
"conversational_response_tool"] |
||||
) |
||||
|
||||
elif self.role == "Editor": |
||||
self.bot: StreamlitBot = EditorBot( |
||||
username=self.username, |
||||
chat=self.chat, |
||||
collection=self.collection, |
||||
project=self.project, |
||||
tools=[ |
||||
"fetch_other_documents_tool", |
||||
"fetch_notes_tool", |
||||
"conversational_response_tool"] |
||||
) |
||||
|
||||
elif self.role == "Podcast": |
||||
st.session_state["make_podcast"] = True |
||||
# with st.sidebar: |
||||
with st.sidebar: |
||||
with st.form("make_podcast_form"): |
||||
instructions = st.text_area( |
||||
"What should the podcast be about? Give a brief description, as if you were the producer." |
||||
) |
||||
start = st.form_submit_button("Make Podcast!") |
||||
if start: |
||||
bot = PodBot( |
||||
subject=self.project.name, |
||||
username=self.username, |
||||
chat=self.chat, |
||||
collection=self.collection, |
||||
project=self.project, |
||||
instructions=instructions |
||||
) |
||||
|
||||
# Save updated chat state to session state |
||||
st.session_state[self.page_name] = { |
||||
"collection": self.collection, |
||||
"project": self.project, |
||||
"chat": self.chat, |
||||
"role": self.role, |
||||
} |
||||
|
||||
# Run the bot (this will display chat history and process user input) |
||||
if self.bot: |
||||
self.bot.run() |
||||
|
||||
else: # If no collection or project is selected, use the conversational response bot |
||||
print_yellow("No collection or project selected. Using conversational response bot.") |
||||
self.bot: StreamlitBot = StreamlitBot( |
||||
username=self.username, |
||||
chat=self.get_chat(), |
||||
tools=["conversational_response_tool"], |
||||
) |
||||
self.bot.run() |
||||
|
||||
|
||||
def get_chat(self, role="Research Assistant", new_chat=False): |
||||
""" |
||||
Retrieves or creates a chat session. |
||||
|
||||
This method handles chat session management by either creating a new chat, |
||||
retrieving an existing one from the database, or initializing a chat when |
||||
none exists in the session state. |
||||
|
||||
Parameters: |
||||
----------- |
||||
role : str, optional |
||||
The role assigned to the chat (default is "Research Assistant"). |
||||
new_chat : bool, optional |
||||
If True, creates a new chat regardless of existing sessions (default is False). |
||||
|
||||
Returns: |
||||
-------- |
||||
StreamlitChat |
||||
A chat instance either newly created or retrieved from the database. |
||||
|
||||
Notes: |
||||
------ |
||||
- If new_chat is True, a new chat is always created |
||||
- If no chat exists in session state, a new one is created |
||||
- Otherwise, retrieves the existing chat from the database using the chat_key in session state |
||||
""" |
||||
print_blue('CHAT TYPE:', role) |
||||
if new_chat: |
||||
chat = StreamlitChat(username=self.username, role=role) |
||||
st.session_state['chat_key'] = chat._key |
||||
print_blue("Creating new chat:", st.session_state['chat_key']) |
||||
elif 'chat_key' not in st.session_state: |
||||
chat = StreamlitChat(username=self.username, role=role) |
||||
st.session_state['chat_key'] = chat._key |
||||
print_blue("Creating new chat:", st.session_state['chat_key']) |
||||
else: |
||||
print_blue("Old chat:", st.session_state['chat_key']) |
||||
chat_data = self.user_arango.db.collection("chats").get(st.session_state['chat_key']) |
||||
chat = StreamlitChat.from_dict(chat_data) |
||||
return chat |
||||
|
||||
def sidebar_actions(self): |
||||
with st.sidebar: |
||||
with st.form("select_chat"): |
||||
self.collection = self.choose_collection("Article collection to use for chat:") |
||||
self.project = self.choose_project("Project to use for chat:") |
||||
submitted = st.form_submit_button("Select Collection/Project") |
||||
|
||||
with st.form("chat_settings"): |
||||
if submitted or any([self.collection, self.project]): |
||||
if self.project: |
||||
self.role = st.selectbox( |
||||
"Choose Bot Role", |
||||
options=["Research Assistant", "Editor", "Podcast"], |
||||
index=0, |
||||
) |
||||
elif self.collection: |
||||
self.role = "Research Assistant" |
||||
|
||||
# Load existing chats from the database |
||||
if self.project: |
||||
chat_history = list( |
||||
self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats FILTER doc["project"] == "{self.project}" RETURN {{"_key": doc["_key"], "name": doc["name"]}}' |
||||
) |
||||
) |
||||
# self.project = Project(username=self.username, project_name=self.project_name, user_arango=self.user_arango) |
||||
elif self.collection: |
||||
chat_history = list( |
||||
self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats FILTER doc["collection"] == "{self.collection}" RETURN {{"_key": doc["_key"], "name": doc["name"]}}' |
||||
) |
||||
) |
||||
|
||||
chats = {i["name"]: i["_key"] for i in chat_history} |
||||
selected_chat = st.selectbox( |
||||
"Continue another chat", options=[""] + list(chats.keys()), index=None |
||||
) |
||||
|
||||
if not self.role: |
||||
self.role == "Research Assistant" |
||||
|
||||
start_chat = st.form_submit_button("Start Chat") |
||||
if start_chat: |
||||
if selected_chat: |
||||
st.session_state["chat_key"] = chats[selected_chat] |
||||
self.chat = self.get_chat() |
||||
else: |
||||
self.chat = self.get_chat(role=self.role, new_chat=True) |
||||
st.rerun() |
||||
|
||||
def remove_old_unsaved_chats(self): |
||||
two_weeks_ago = datetime.now() - timedelta(weeks=2) |
||||
q = f'FOR doc IN chats FILTER doc.saved == false AND doc.last_updated < "{two_weeks_ago.isoformat()}" RETURN doc' |
||||
print_blue(q) |
||||
old_chats = self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats RETURN doc' |
||||
) |
||||
print('test', old_chats) |
||||
old_chats = self.user_arango.db.aql.execute( |
||||
f'FOR doc IN chats FILTER doc.saved == false AND doc.last_updated < "{two_weeks_ago.isoformat()}" RETURN doc' |
||||
) |
||||
for chat in old_chats: |
||||
print_red(chat["_id"]) |
||||
self.user_arango.db.collection("chats").delete(chat["_key"]) |
||||
|
||||
@ -0,0 +1,103 @@ |
||||
from _rss import RSSReader |
||||
import streamlit as st |
||||
from _base_class import StreamlitBaseClass |
||||
from colorprinter.print_color import * |
||||
|
||||
class RSSFeedsPage(StreamlitBaseClass): |
||||
def __init__(self, username: str): |
||||
super().__init__(username=username) |
||||
self.page_name = "RSS Feeds" |
||||
self.reader = RSSReader(username=username) |
||||
# Initialize attributes from session state if available |
||||
for k, v in st.session_state.get(self.page_name, {}).items(): |
||||
setattr(self, k, v) |
||||
|
||||
def run(self): |
||||
if "selected_feed" not in st.session_state: |
||||
st.session_state["selected_feed"] = None |
||||
self.update_current_page(self.page_name) |
||||
self.display_feed() |
||||
self.sidebar_actions() |
||||
self.update_session_state(page_name=self.page_name) |
||||
|
||||
def select_rss_feeds(self): |
||||
rss_feeds = self.reader.get_rss_feeds() |
||||
if rss_feeds: |
||||
feed_options = [feed["title"] for feed in rss_feeds] |
||||
with st.sidebar: |
||||
st.subheader("Show your feeds") |
||||
selected_feed_title = st.selectbox( |
||||
"Select a feed", options=feed_options, index=None |
||||
) |
||||
if selected_feed_title: |
||||
st.session_state["selected_feed"] = [ |
||||
feed["_key"] |
||||
for feed in rss_feeds |
||||
if feed["title"] == selected_feed_title |
||||
][0] |
||||
st.rerun() |
||||
else: |
||||
st.write("You have no RSS feeds added.") |
||||
|
||||
def search_feeds(self, rss_url): |
||||
with st.spinner("Discovering feeds..."): |
||||
feeds = self.reader.discover_feeds(rss_url) |
||||
if feeds: |
||||
st.session_state["discovered_feeds"] = feeds |
||||
else: |
||||
st.error("No RSS feeds found at the provided URL.") |
||||
|
||||
def sidebar_actions(self): |
||||
if "discovered_feeds" not in st.session_state: |
||||
st.session_state["discovered_feeds"] = None |
||||
|
||||
with st.sidebar: |
||||
self.select_rss_feeds() |
||||
st.subheader("Add a New RSS Feed") |
||||
with st.form("add_rss_feed"): |
||||
rss_url = st.text_input("Website URL or RSS Feed URL") |
||||
submitted = st.form_submit_button("Discover Feeds") |
||||
if submitted: |
||||
print_green(rss_url) |
||||
feeds = self.reader.discover_feeds(rss_url) |
||||
st.session_state["discovered_feeds"] = feeds |
||||
|
||||
if st.session_state["discovered_feeds"]: |
||||
st.subheader("Select a Feed to Add") |
||||
feeds = st.session_state["discovered_feeds"] |
||||
feed_options = [f"{feed['title']} ({feed['href']})" for feed in feeds] |
||||
selected_feed = st.selectbox("Available Feeds", options=feed_options) |
||||
selected_feed_url = feeds[feed_options.index(selected_feed)]["href"] |
||||
|
||||
if st.button("Preview Feed"): |
||||
feed = self.reader.parse_feed(selected_feed_url) |
||||
st.write(f"{feed.title}") |
||||
description = self.reader.html_to_markdown(feed.description) |
||||
st.write(f"_{description}_") |
||||
for entry in feed.entries[:5]: |
||||
with st.expander(entry["title"]): |
||||
summary = entry.get("summary", "No summary available") |
||||
markdown_summary = self.reader.html_to_markdown(summary) |
||||
st.markdown(markdown_summary) |
||||
print_yellow(selected_feed_url) |
||||
|
||||
if st.button( |
||||
"Add RSS Feed", |
||||
on_click=self.reader.add_rss_feed, |
||||
args=[selected_feed_url], |
||||
): |
||||
del st.session_state["discovered_feeds"] |
||||
st.success("RSS Feed added.") |
||||
st.rerun() |
||||
|
||||
def display_feed(self): |
||||
if st.session_state["selected_feed"]: |
||||
self.reader.get_feed(st.session_state["selected_feed"]) |
||||
st.title(self.reader.feed.title) |
||||
st.write(f"_{self.reader.feed.description}_") |
||||
for entry in self.reader.feed.entries[:5]: |
||||
with st.expander(entry["title"]): |
||||
summary = entry.get("summary", "No summary available") |
||||
markdown_summary = self.reader.html_to_markdown(summary) |
||||
st.markdown(markdown_summary) |
||||
st.markdown(f"[Read more]({entry['link']})") |
||||
@ -0,0 +1,6 @@ |
||||
from pydantic import BaseModel |
||||
|
||||
class QueryResponse(BaseModel): |
||||
query_to_vector_database: str |
||||
short_explanation: str |
||||
|
||||
@ -0,0 +1,406 @@ |
||||
import streamlit as st |
||||
from datetime import datetime |
||||
from colorprinter.print_color import * |
||||
|
||||
from _base_class import StreamlitBaseClass |
||||
from projects_page import Project |
||||
from agent_research import ResearchReport, MasterAgent, StructureAgent, ToolAgent, ArchiveAgent, process_step |
||||
import os |
||||
import json |
||||
|
||||
|
||||
class ResearchPage(StreamlitBaseClass): |
||||
""" |
||||
ResearchPage - A Streamlit interface for deep research using AI agents. |
||||
|
||||
This class provides a user interface for conducting in-depth research using |
||||
multiple specialized AI agents working together. It allows users to input |
||||
research questions, track progress, and view detailed research reports. |
||||
|
||||
Attributes: |
||||
username (str): The username of the current user. |
||||
project_name (str): Name of the selected project. |
||||
project (Project): Project instance the research is associated with. |
||||
page_name (str): Name of the current page ("Research"). |
||||
research_state (dict): Dictionary tracking the current state of research. |
||||
report (ResearchReport): Instance for tracking research progress and results. |
||||
|
||||
Methods: |
||||
run(): Main method to render the research interface and handle interactions. |
||||
sidebar_actions(): Renders sidebar elements for selecting projects and research options. |
||||
start_new_research(): Initiates a new research session. |
||||
view_saved_reports(): Displays a list of saved research reports. |
||||
display_report(): Renders a research report in the Streamlit interface. |
||||
show_research_progress(): Displays the current research progress. |
||||
""" |
||||
def __init__(self, username): |
||||
super().__init__(username=username) |
||||
self.project_name = None |
||||
self.project = None |
||||
self.page_name = "Research" |
||||
|
||||
# Research state tracking |
||||
self.research_state = { |
||||
"in_progress": False, |
||||
"completed": False, |
||||
"question": None, |
||||
"started_at": None, |
||||
"report": None, |
||||
"current_step": None, |
||||
"steps_completed": 0, |
||||
"total_steps": 0 |
||||
} |
||||
|
||||
self.report = None |
||||
|
||||
# Initialize attributes from session state if available |
||||
if self.page_name in st.session_state: |
||||
for k, v in st.session_state[self.page_name].items(): |
||||
setattr(self, k, v) |
||||
|
||||
# Create reports directory if it doesn't exist |
||||
os.makedirs(f"/home/lasse/sci/reports", exist_ok=True) |
||||
|
||||
def run(self): |
||||
self.update_current_page("Research") |
||||
self.sidebar_actions() |
||||
|
||||
st.title("Deep Research") |
||||
|
||||
if not self.project: |
||||
st.warning("Please select a project to start researching.") |
||||
return |
||||
|
||||
# Main interface |
||||
if self.research_state["in_progress"]: |
||||
self.show_research_progress() |
||||
elif self.research_state["completed"]: |
||||
self.display_report(self.research_state["report"]) |
||||
else: |
||||
# Input for new research |
||||
st.subheader(f"New Research for Project: {self.project_name}") |
||||
with st.form("research_form"): |
||||
question = st.text_area("Enter your research question:", |
||||
help="Be specific about what you want to research. Complex questions will be broken down into sub-questions.") |
||||
start_button = st.form_submit_button("Start Research") |
||||
|
||||
if start_button and question: |
||||
self.start_new_research(question) |
||||
st.rerun() |
||||
|
||||
# Option to view saved reports |
||||
with st.expander("View Saved Reports"): |
||||
self.view_saved_reports() |
||||
|
||||
def sidebar_actions(self): |
||||
with st.sidebar: |
||||
with st.form("select_project"): |
||||
self.project = self.choose_project("Project for research:") |
||||
submitted = st.form_submit_button("Select Project") |
||||
|
||||
if submitted and self.project: |
||||
self.project_name = self.project.name |
||||
st.success(f"Selected project: {self.project_name}") |
||||
|
||||
if self.research_state["in_progress"]: |
||||
st.info(f"Research in progress: {self.research_state['question']}") |
||||
if st.button("Cancel Research"): |
||||
self.research_state["in_progress"] = False |
||||
st.rerun() |
||||
|
||||
elif self.research_state["completed"]: |
||||
if st.button("Start New Research"): |
||||
self.research_state["completed"] = False |
||||
self.research_state["report"] = None |
||||
st.rerun() |
||||
|
||||
def start_new_research(self, question): |
||||
"""Initiates a new research session with the given question""" |
||||
self.research_state["question"] = question |
||||
self.research_state["in_progress"] = True |
||||
self.research_state["completed"] = False |
||||
self.research_state["started_at"] = datetime.now().isoformat() |
||||
|
||||
# Initialize the research report |
||||
self.report = ResearchReport( |
||||
question=question, |
||||
username=self.username, |
||||
project_name=self.project_name |
||||
) |
||||
|
||||
# Save current state |
||||
st.session_state[self.page_name] = { |
||||
"project_name": self.project_name, |
||||
"project": self.project, |
||||
"research_state": self.research_state, |
||||
"report": self.report |
||||
} |
||||
|
||||
# Start a new thread to run the research process |
||||
# In a production environment, you might want to use a background job |
||||
# For now, we'll run it in the main thread with streamlit spinner |
||||
with st.spinner("Research in progress... This may take several minutes."): |
||||
try: |
||||
# Initialize agents |
||||
master_agent = MasterAgent( |
||||
username=self.username, |
||||
project=self.project, |
||||
report=self.report, |
||||
chat=True |
||||
) |
||||
structure_agent = StructureAgent( |
||||
username=self.username, |
||||
model="small", |
||||
report=self.report |
||||
) |
||||
tool_agent = ToolAgent( |
||||
username=self.username, |
||||
model="tools", |
||||
system_message="You are an assistant with tools. Always choose a tool to help with the task.", |
||||
report=self.report, |
||||
project=self.project, |
||||
chat=True |
||||
) |
||||
archive_agent = ArchiveAgent( |
||||
username=self.username, |
||||
report=self.report, |
||||
project=self.project, |
||||
system_message="You are an assistant specialized in reading and summarizing research information.", |
||||
chat=True |
||||
) |
||||
|
||||
# Track the research state in the master agent |
||||
master_agent.research_state["original_question"] = question |
||||
|
||||
# Execute the research workflow |
||||
# 1. Create research plan |
||||
st.text("Creating research plan...") |
||||
research_plan = master_agent.make_plan(question) |
||||
self.report.log_plan(research_plan) |
||||
|
||||
# 2. Structure the plan |
||||
st.text("Structuring research plan...") |
||||
structured_plan = structure_agent.make_structured(research_plan, question) |
||||
self.report.log_plan(research_plan, structured_plan.model_dump()) |
||||
|
||||
# Update total steps count |
||||
self.research_state["total_steps"] = len(structured_plan.steps) |
||||
|
||||
# 3. Execute the plan step by step |
||||
execution_results = {} |
||||
|
||||
for step_name, tasks in structured_plan.steps.items(): |
||||
st.text(f"Processing step: {step_name}") |
||||
self.research_state["current_step"] = step_name |
||||
self.research_state["steps_completed"] += 1 |
||||
|
||||
# Collect all task descriptions in this step |
||||
step_tasks = [ |
||||
{"task_name": task_name, "task_description": task_description} |
||||
for task_name, task_description in tasks |
||||
] |
||||
|
||||
# Process the entire step |
||||
step_result = process_step( |
||||
step_name, step_tasks, master_agent, tool_agent, archive_agent |
||||
) |
||||
execution_results[step_name] = step_result |
||||
|
||||
# 4. Evaluate if more steps are needed |
||||
st.text("Evaluating research plan...") |
||||
plan_evaluation = master_agent.evaluate_plan(execution_results) |
||||
self.report.log_plan_evaluation(plan_evaluation) |
||||
|
||||
# 5. Write the final report |
||||
st.text("Writing final report...") |
||||
final_report = master_agent.write_report(execution_results) |
||||
self.report.log_final_report(final_report) |
||||
|
||||
# 6. Save the reports |
||||
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") |
||||
report_path = f"/home/lasse/sci/reports/research_report_{self.username}_{timestamp}" |
||||
|
||||
# Save JSON report |
||||
json_path = f"{report_path}.json" |
||||
with open(json_path, "w") as f: |
||||
json.dump(self.report.get_full_report(), f, indent=2) |
||||
|
||||
# Save markdown report |
||||
markdown_report = self.report.get_markdown_report() |
||||
markdown_path = f"{report_path}.md" |
||||
with open(markdown_path, "w") as f: |
||||
f.write(markdown_report) |
||||
|
||||
# Update research state |
||||
self.research_state["in_progress"] = False |
||||
self.research_state["completed"] = True |
||||
self.research_state["report"] = { |
||||
"json_path": json_path, |
||||
"markdown_path": markdown_path, |
||||
"report_data": self.report.get_full_report(), |
||||
"markdown_content": markdown_report |
||||
} |
||||
|
||||
except Exception as e: |
||||
st.error(f"An error occurred during research: {str(e)}") |
||||
import traceback |
||||
st.code(traceback.format_exc()) |
||||
self.research_state["in_progress"] = False |
||||
|
||||
# Update session state |
||||
st.session_state[self.page_name] = { |
||||
"project_name": self.project_name, |
||||
"project": self.project, |
||||
"research_state": self.research_state, |
||||
"report": self.report |
||||
} |
||||
|
||||
def view_saved_reports(self): |
||||
"""Displays a list of saved research reports""" |
||||
reports_dir = "/home/lasse/sci/reports" |
||||
if not os.path.exists(reports_dir): |
||||
st.info("No saved reports found.") |
||||
return |
||||
|
||||
# Get all report files |
||||
json_files = [f for f in os.listdir(reports_dir) if f.endswith('.json') and f.startswith('research_report')] |
||||
|
||||
if not json_files: |
||||
st.info("No saved reports found.") |
||||
return |
||||
|
||||
for file in sorted(json_files, reverse=True): |
||||
file_path = os.path.join(reports_dir, file) |
||||
try: |
||||
with open(file_path, 'r') as f: |
||||
report_data = json.load(f) |
||||
|
||||
# Extract basic info |
||||
question = report_data.get("metadata", {}).get("question", "Unknown question") |
||||
project = report_data.get("metadata", {}).get("project_name", "No project") |
||||
started_at = report_data.get("metadata", {}).get("started_at", "Unknown time") |
||||
|
||||
# Format the date |
||||
try: |
||||
date_obj = datetime.fromisoformat(started_at) |
||||
date_str = date_obj.strftime("%Y-%m-%d %H:%M") |
||||
except: |
||||
date_str = started_at |
||||
|
||||
# Create an expandable section for each report |
||||
st.markdown(f"_{question} ({project} - {date_str})_") |
||||
st.markdown(f"**Project:** {project}") |
||||
st.markdown(f"**Date:** {date_str}") |
||||
|
||||
# Button to view full report |
||||
if st.button("View Full Report", key=f"view_{file}"): |
||||
# Load corresponding markdown file if it exists |
||||
md_file = file.replace('.json', '.md') |
||||
md_path = os.path.join(reports_dir, md_file) |
||||
|
||||
if os.path.exists(md_path): |
||||
with open(md_path, 'r') as f: |
||||
markdown_content = f.read() |
||||
else: |
||||
markdown_content = None |
||||
|
||||
self.research_state["completed"] = True |
||||
self.research_state["report"] = { |
||||
"json_path": file_path, |
||||
"markdown_path": md_path if os.path.exists(md_path) else None, |
||||
"report_data": report_data, |
||||
"markdown_content": markdown_content |
||||
} |
||||
st.rerun() |
||||
|
||||
except Exception as e: |
||||
st.error(f"Error loading report {file}: {str(e)}") |
||||
|
||||
def display_report(self, report_data): |
||||
"""Renders a research report in the Streamlit interface""" |
||||
if not report_data: |
||||
st.warning("No report data available.") |
||||
return |
||||
|
||||
st.title("Research Report") |
||||
|
||||
# Get report data |
||||
markdown_content = report_data.get("markdown_content") |
||||
json_data = report_data.get("report_data") |
||||
|
||||
if markdown_content: |
||||
# Display the markdown report |
||||
st.markdown(markdown_content) |
||||
elif json_data: |
||||
# Fallback to displaying JSON data in a more readable format |
||||
question = json_data.get("metadata", {}).get("question", "Unknown question") |
||||
st.header(f"Research on: {question}") |
||||
|
||||
# Display metadata |
||||
st.subheader("Metadata") |
||||
metadata = json_data.get("metadata", {}) |
||||
st.markdown(f"**Project:** {metadata.get('project_name', 'None')}") |
||||
st.markdown(f"**Started:** {metadata.get('started_at', 'Unknown')}") |
||||
st.markdown(f"**Finished:** {metadata.get('finished_at', 'Unknown')}") |
||||
|
||||
# Display final report |
||||
st.subheader("Research Findings") |
||||
st.markdown(json_data.get("final_report", "No final report available.")) |
||||
|
||||
# Display steps |
||||
st.subheader("Research Steps") |
||||
steps = json_data.get("steps", {}) |
||||
for step_name, step_data in steps.items(): |
||||
with st.expander(step_name): |
||||
st.markdown(f"**Summary:** {step_data.get('summary', 'No summary available.')}") |
||||
|
||||
# Display tools used |
||||
st.markdown("**Tools used:**") |
||||
for tool in step_data.get("tools_used", []): |
||||
st.markdown(f"- {tool.get('tool', 'Unknown tool')} with query: _{tool.get('args', {}).get('query', 'No query')}_") |
||||
|
||||
else: |
||||
st.error("No report content available to display.") |
||||
|
||||
# Download buttons |
||||
col1, col2 = st.columns(2) |
||||
with col1: |
||||
if report_data.get("markdown_path") and os.path.exists(report_data["markdown_path"]): |
||||
with open(report_data["markdown_path"], "r") as f: |
||||
markdown_content = f.read() |
||||
st.download_button( |
||||
label="Download as Markdown", |
||||
data=markdown_content, |
||||
file_name=os.path.basename(report_data["markdown_path"]), |
||||
mime="text/markdown" |
||||
) |
||||
|
||||
with col2: |
||||
if report_data.get("json_path") and os.path.exists(report_data["json_path"]): |
||||
with open(report_data["json_path"], "r") as f: |
||||
json_content = f.read() |
||||
st.download_button( |
||||
label="Download as JSON", |
||||
data=json_content, |
||||
file_name=os.path.basename(report_data["json_path"]), |
||||
mime="application/json" |
||||
) |
||||
|
||||
def show_research_progress(self): |
||||
"""Displays the current research progress""" |
||||
st.subheader("Research in Progress") |
||||
st.markdown(f"**Question:** {self.research_state['question']}") |
||||
|
||||
# Show progress bar |
||||
progress = 0 |
||||
if self.research_state["total_steps"] > 0: |
||||
progress = self.research_state["steps_completed"] / self.research_state["total_steps"] |
||||
|
||||
st.progress(progress) |
||||
|
||||
# Show current step |
||||
current_step = self.research_state.get("current_step", "Planning") |
||||
st.markdown(f"**Current step:** {current_step}") |
||||
|
||||
st.info("Research is ongoing. This may take several minutes depending on the complexity of the question.") |
||||
st.warning("Please do not navigate away from this page while research is in progress.") |
||||
@ -0,0 +1,55 @@ |
||||
|
||||
import streamlit as st |
||||
from time import sleep |
||||
from colorprinter.print_color import * |
||||
|
||||
|
||||
from _base_class import StreamlitBaseClass |
||||
|
||||
|
||||
|
||||
class SettingsPage(StreamlitBaseClass): |
||||
def __init__(self, username: str): |
||||
super().__init__(username=username) |
||||
|
||||
def run(self): |
||||
self.update_current_page("Settings") |
||||
self.set_profile_picture() |
||||
self.use_reasoning_model() |
||||
|
||||
def set_profile_picture(self): |
||||
st.markdown("Profile picture") |
||||
profile_picture = st.file_uploader( |
||||
"Upload profile picture", type=["png", "jpg", "jpeg"] |
||||
) |
||||
if profile_picture: |
||||
# Resize the image to 64x64 pixels |
||||
from PIL import Image |
||||
|
||||
img = Image.open(profile_picture) |
||||
img.thumbnail((64, 64)) |
||||
img_path = f"user_data/{st.session_state['username']}/profile_picture.png" |
||||
img.save(img_path) |
||||
self.update_settings("avatar", img_path) |
||||
st.success("Profile picture uploaded") |
||||
sleep(1) |
||||
|
||||
def use_reasoning_model(self): |
||||
""" |
||||
Displays a checkbox in the Streamlit interface to enable or disable the reasoning model for generating responses in chats. |
||||
|
||||
Retrieves the current settings and checks if the "use_reasoning_model" key exists. If not, it initializes it to False. |
||||
Then, it displays a markdown text and a checkbox for the user to toggle the reasoning model usage. |
||||
The updated setting is saved back to the settings. |
||||
|
||||
Returns: |
||||
None |
||||
""" |
||||
settings = self.get_settings() |
||||
if "use_reasoning_model" not in settings: |
||||
settings["use_reasoning_model"] = False |
||||
st.markdown("Use Reasoning Model") |
||||
|
||||
use_reasoning_model = st.checkbox("Use Reasoning Model", value=settings["use_reasoning_model"], help="Use the reasoning model to generate responses in chats. This may take longer to process.") |
||||
self.update_settings("use_reasoning_model", use_reasoning_model) |
||||
|
||||
@ -0,0 +1,206 @@ |
||||
from _llm import LLM |
||||
from _arango import ArangoDB |
||||
from _chromadb import ChromaDB |
||||
from streamlit_chatbot import Bot |
||||
from pydantic import BaseModel, Field |
||||
from typing import Dict, List, Tuple |
||||
from colorprinter.print_color import * |
||||
from projects_page import Project |
||||
from _base_class import StreamlitBaseClass |
||||
from prompts import get_tools_prompt |
||||
|
||||
class ResearchBase(Bot): |
||||
def __init__(self, username, **args): |
||||
super().__init__(username=username, **args) |
||||
self.llm = LLM() |
||||
self.arango = ArangoDB() |
||||
self.chromadb = ChromaDB() |
||||
self.messages = [] |
||||
|
||||
def start(self): |
||||
self.messages = [{"role": "system", "message": self.llm.system_message}] |
||||
if self.llm.model in ["small", "standard", "vision", "reasoning", "tools"]: |
||||
self.llm.get_model(self.llm.model) |
||||
|
||||
|
||||
class ResearchManager(ResearchBase): |
||||
def __init__(self, username, project=None): |
||||
super().__init__(username=username, project=project) |
||||
self.llm.system_message = "You are an assistant helping a journalist writing a report based on extensive research." |
||||
self.llm.model = "reasoning" |
||||
self.start() |
||||
|
||||
def generate_plan(self, question): |
||||
query = f""" |
||||
A journalist wants to get a report that answers this question: "{question}" |
||||
THIS IS *NOT* A QUESTION YOU CAN ANSWER! Instead, you need to make a plan for how to answer this question. |
||||
Include what type of information you need from what available sources. |
||||
Available sources are: |
||||
- Scientific articles |
||||
- Other articles the journalists has gathered, such as blog posts, news articles, etc. |
||||
- The journalists own notes. |
||||
- Transcribed interviews (already done, you can't produce new ones). |
||||
All of the above sources are available in a database, but you need to specify what you need. Be as precise as possible. |
||||
As you don't have access to the sources, include steps to retrieve excerpts from articles and retrieve those that might be interesting. |
||||
Also include steps to verify the information. |
||||
Make the plan easy to follow and structured. |
||||
Remember: You are not answering the question, you are making *a plan* for how to answer the question using the available sources. |
||||
""" |
||||
query += f"\nTo help you understand the subject, here is a summary of notes the journalist has done: {project.notes_summary}" |
||||
query += """Please structure the plan like: |
||||
## Step 1: |
||||
- Task1: Description of task |
||||
- Task2: Description of task |
||||
## Step 2: |
||||
- Task1: Description of task |
||||
- Task2: Description of task |
||||
Etc, with as many steps and tasks as needed. |
||||
""" |
||||
return self.llm.generate(query).content |
||||
|
||||
|
||||
class ResearchAssistant(ResearchBase): |
||||
def __init__(self, username): |
||||
super().__init__(username) |
||||
self.llm.system_message = "You are a Research Assistant" |
||||
self.start() |
||||
|
||||
|
||||
class HelperBot(ResearchBase): |
||||
def __init__(self, username): |
||||
super().__init__(username) |
||||
self.llm.system_message = "You are helping a researcher to structure a text. You will get a text and make it into structured data. Make sure not to change the meaning of the text and keeps all the details in the subtasks." |
||||
self.llm.model = "small" |
||||
self.start() |
||||
|
||||
def make_structured_plan(self, text, question=None): |
||||
|
||||
class Plan(BaseModel): |
||||
steps: Dict[str, List[Tuple[str, str]]] = Field( |
||||
description="Structured plan represented as steps with their corresponding tasks or facts", |
||||
example={ |
||||
"Step 1: Gather Existing Materials": [ |
||||
("Task 1", "Description of task"), |
||||
("Task 2", "Description of task"), |
||||
], |
||||
"Step 2: Extract Relevant Information": [ |
||||
("Task 1", "Description of task"), |
||||
("Task 2", "Description of task"), |
||||
], |
||||
}, |
||||
) |
||||
|
||||
if question: |
||||
query = f''' This is a proposed plan for how to write a report on "{question}":\n"""{text}"""\nPlease make the plan into structured data with subtasks. Make sure to keep all the details in the subtasks.''' |
||||
else: |
||||
query = f''' This is a proposed plan for how to write a report:\n"""{text}"""\nPlease make the plan into structured data with subtasks. Make sure to keep all the details in the subtasks.''' |
||||
response = self.llm.generate(query, format=Plan.model_json_schema()) |
||||
print(response) |
||||
structured_response = Plan.model_validate_json(response.content) |
||||
print('PLAN') |
||||
print_rainbow(structured_response) |
||||
print() |
||||
return structured_response |
||||
|
||||
|
||||
class ToolBot(ResearchBase): |
||||
def __init__(self, username, tools: list): |
||||
super().__init__(username, tools=tools) |
||||
self.start() |
||||
tools_names = [tool.__name__ for tool in self.tools] |
||||
tools_name_string = "\n– ".join(tools_names) |
||||
self.llm = LLM( |
||||
temperature=0, |
||||
system_message=f""" |
||||
You are an helpful assistant with tools. The tools you can choose from are: |
||||
{tools_name_string} |
||||
Your task is to choose one or multiple tools to answering a user's query. |
||||
DON'T come up with your own tools, only use the ones provided. |
||||
""", |
||||
chat=False, |
||||
model="tools", |
||||
) |
||||
|
||||
def propose_tools(self, task): |
||||
query = f"""What tool(s) would you use to help with this task: |
||||
"{task}" |
||||
Answer in a structured way using the tool_calls field! |
||||
""" |
||||
query = get_tools_prompt(task) |
||||
response = self.llm.generate(query) |
||||
print_yellow('Model:', self.llm.model) |
||||
print_rainbow(response) |
||||
return response.tool_calls |
||||
|
||||
if __name__ == "__main__": |
||||
|
||||
base = StreamlitBaseClass(username="lasse") |
||||
project = Project( |
||||
username="lasse", |
||||
project_name="Monarch butterflies", |
||||
user_arango=base.get_arango(), |
||||
) |
||||
rm = ResearchManager(username="lasse", project=project) |
||||
tb = ToolBot( |
||||
username="lasse", |
||||
tools=[ |
||||
"fetch_science_articles_tool", |
||||
"fetch_notes_tool", |
||||
"fetch_other_documents_tool", |
||||
"fetch_science_articles_and_other_documents_tool", |
||||
] |
||||
) |
||||
# ra = ResearchAssistant(username="lasse") |
||||
hb = HelperBot(username="lasse") |
||||
|
||||
question = "Tell me five interesting facts about the Monarch butterfly" |
||||
|
||||
# Generate plan |
||||
plan = rm.generate_plan(question) |
||||
# -- Example of what a plan can look like -- |
||||
# plan = """## Step-by-Step Plan for Answering the Question: "Tell Me Five Interesting Facts About the Monarch Butterfly" |
||||
|
||||
# ### Step 1: Gather and Organize Existing Materials |
||||
# - **Task 1:** Retrieve all existing materials related to Monarch butterflies from the database using keywords such as "Monarch butterfly migration," "habitat loss," "milkweed," "insecticides," "climate change," "Monarch Butterfly Biosphere Reserve," and "migration patterns." |
||||
# - **Task 2:** Categorize these materials into scientific articles, other articles (blogs, news), own notes, and transcribed interviews for easy access. |
||||
|
||||
# ### Step 2: Extract Relevant Excerpts |
||||
# - **Task 1:** From the retrieved scientific articles, extract information on migration patterns, genetic studies, and population decline factors. |
||||
# - **Task 2:** From blogs and news articles, look for interesting anecdotes or recent findings about conservation efforts and unique behaviors of Monarch butterflies. |
||||
|
||||
# ### Step 3: Identify Potential Interesting Facts |
||||
# - **Task 1:** Review the extracted excerpts to identify potential facts such as migration patterns, threats faced by Monarchs, population decline statistics, conservation efforts, and unique behaviors. |
||||
# - **Task 2:** Compile a list of five compelling and accurate facts based on the extracted information. |
||||
|
||||
# ### Step 4: Verify Information |
||||
# - **Task 1:** Cross-check each fact with multiple sources to ensure accuracy. For example, verify migration details across scientific articles and recent news reports. |
||||
# - **Task 2:** Look for consensus among sources regarding population trends and threats to Monarchs. |
||||
|
||||
# ### Step 5: Structure the Report |
||||
# - **Task 1:** Organize the five selected facts into a coherent structure, ensuring each fact is clearly explained and engaging. |
||||
# - **Task 2:** Incorporate quotes or statistics from sources to add depth and credibility to each fact. |
||||
|
||||
# ### Step 6: Review and Finalize |
||||
# - **Task 1:** Proofread the report for clarity, accuracy, and grammar. |
||||
# - **Task 2:** Ensure all information is presented in an engaging manner suitable for a journalistic report. |
||||
|
||||
# This plan ensures that the journalist systematically gathers, verifies, and presents five interesting facts about Monarch butterflies, providing a comprehensive and accurate report. |
||||
# """ |
||||
#print_blue(plan) |
||||
if "</think>" in plan: |
||||
plan = plan.split("</think>")[1] |
||||
|
||||
# Make structured plan |
||||
structured_plan = hb.make_structured_plan(plan, question) |
||||
|
||||
|
||||
for step, tasks in structured_plan.steps.items(): |
||||
print_blue("\n### Step:", step) |
||||
for task in tasks: |
||||
|
||||
print_blue("Task:", task[0]) |
||||
print_yellow(task[1]) |
||||
|
||||
tools = tb.propose_tools(task[1]) |
||||
print_green("Tools:", tools) |
||||
print('\n') |
||||
Loading…
Reference in new issue