|
|
from datetime import datetime |
|
|
import streamlit as st |
|
|
import uuid |
|
|
|
|
|
from _base_class import StreamlitBaseClass, BaseClass |
|
|
from _llm import LLM |
|
|
from _arango import ArangoDB |
|
|
from prompts import * |
|
|
from colorprinter.print_color import * |
|
|
from llm_tools import ToolRegistry |
|
|
from streamlit_chatbot import StreamlitBot, PodBot, EditorBot, ResearchAssistantBot |
|
|
|
|
|
class Chat(StreamlitBaseClass): |
|
|
def __init__(self, username=None, **kwargs): |
|
|
super().__init__(username=username, **kwargs) |
|
|
self.name = kwargs.get("name", None) |
|
|
self.chat_history = kwargs.get("chat_history", []) |
|
|
self.role = kwargs.get("role", "Research Assistant") |
|
|
self._key = kwargs.get("_key", str(uuid.uuid4())) |
|
|
self.saved = kwargs.get("saved", False) |
|
|
self.last_updated = kwargs.get("last_updated", datetime.now().isoformat()) |
|
|
self.message_attachments = None |
|
|
self.project = kwargs.get("project", None) |
|
|
|
|
|
def add_message(self, role, content): |
|
|
self.chat_history.append( |
|
|
{ |
|
|
"role": role, |
|
|
"content": content.strip().strip('"'), |
|
|
"role_type": self.role, |
|
|
} |
|
|
) |
|
|
|
|
|
def to_dict(self): |
|
|
return { |
|
|
"_key": self._key, |
|
|
"name": self.name, |
|
|
"chat_history": self.chat_history, |
|
|
"role": self.role, |
|
|
"username": self.username, |
|
|
"project": self.project, |
|
|
"last_updated": self.last_updated, |
|
|
"saved": self.saved, |
|
|
} |
|
|
|
|
|
def update_in_arango(self): |
|
|
"""Update chat in ArangoDB using the new API""" |
|
|
self.last_updated = datetime.now().isoformat() |
|
|
|
|
|
# Use the create_or_update_chat method from the new API |
|
|
self.user_arango.create_or_update_chat(self.to_dict()) |
|
|
|
|
|
def set_name(self, user_input): |
|
|
llm = LLM( |
|
|
model="small", |
|
|
max_length_answer=50, |
|
|
temperature=0.4, |
|
|
system_message="You are a chatbot who will be chatting with a user", |
|
|
) |
|
|
prompt = ( |
|
|
f'Give a short name to the chat based on this user input: "{user_input}" ' |
|
|
"No more than 30 characters. Answer ONLY with the name of the chat." |
|
|
) |
|
|
name = llm.generate(prompt).content.strip('"') |
|
|
name = f'{name} - {datetime.now().strftime("%B %d")}' |
|
|
|
|
|
# Check for existing chat with the same name |
|
|
existing_chat = self.user_arango.execute_aql( |
|
|
""" |
|
|
FOR chat IN chats |
|
|
FILTER chat.name == @name AND chat.username == @username |
|
|
RETURN chat |
|
|
""", |
|
|
bind_vars={"name": name, "username": self.username} |
|
|
) |
|
|
|
|
|
if list(existing_chat): |
|
|
name = f'{name} ({datetime.now().strftime("%H:%M")})' |
|
|
name += f" - [{self.role}]" |
|
|
self.name = name |
|
|
return name |
|
|
|
|
|
def show_chat_history(self): |
|
|
"""Display chat history in the Streamlit UI""" |
|
|
for message in self.chat_history: |
|
|
with st.chat_message( |
|
|
name="assistant" if message["role"] == "assistant" else "user", |
|
|
avatar=self.get_avatar(role=message["role"]) |
|
|
): |
|
|
st.write(message["content"]) |
|
|
|
|
|
def get_avatar(self, role): |
|
|
"""Get avatar for a role""" |
|
|
if role == "user": |
|
|
return None |
|
|
elif role == "Host": |
|
|
return "🎙️" |
|
|
elif role == "Guest": |
|
|
return "🎤" |
|
|
elif role == "assistant": |
|
|
if self.role == "Research Assistant": |
|
|
return "🔬" |
|
|
elif self.role == "Editor": |
|
|
return "📝" |
|
|
else: |
|
|
return "🤖" |
|
|
return None |
|
|
|
|
|
@classmethod |
|
|
def from_dict(cls, data): |
|
|
return cls( |
|
|
username=data.get("username"), |
|
|
name=data.get("name"), |
|
|
chat_history=data.get("chat_history", []), |
|
|
role=data.get("role", "Research Assistant"), |
|
|
_key=data.get("_key"), |
|
|
project=data.get("project"), |
|
|
last_updated=data.get("last_updated"), |
|
|
saved=data.get("saved", False), |
|
|
) |
|
|
|
|
|
def chat_history2bot(self, n_messages: int = None, remove_system: bool = False): |
|
|
history = [ |
|
|
{"role": m["role"], "content": m["content"]} for m in self.chat_history |
|
|
] |
|
|
if n_messages and len(history) > n_messages: |
|
|
history = history[-n_messages:] |
|
|
if ( |
|
|
all([history[0]["role"] == "system", remove_system]) |
|
|
or history[0]["role"] == "assistant" |
|
|
): |
|
|
history = history[1:] |
|
|
return history |
|
|
|
|
|
|
|
|
class Bot(BaseClass): |
|
|
def __init__(self, username: str, chat: Chat = None, tools: list = None, **kwargs): |
|
|
super().__init__(username=username, **kwargs) |
|
|
|
|
|
# Use the passed in chat or create a new Chat |
|
|
self.chat = chat if chat else Chat(username=username, role="Research Assistant") |
|
|
print_yellow(f"Chat:", chat, type(chat)) |
|
|
|
|
|
# Store or set up project/collection if available |
|
|
self.project = kwargs.get("project", None) |
|
|
self.collection = kwargs.get("collection", None) |
|
|
if self.collection and not isinstance(self.collection, list): |
|
|
self.collection = [self.collection] |
|
|
|
|
|
# Load articles in the collections using the new API |
|
|
self.arango_ids = [] |
|
|
if self.collection: |
|
|
for c in self.collection: |
|
|
# Use execute_aql from the new API |
|
|
article_ids = self.user_arango.execute_aql( |
|
|
""" |
|
|
FOR doc IN article_collections |
|
|
FILTER doc.name == @collection |
|
|
FOR article IN doc.articles |
|
|
RETURN article |
|
|
""", |
|
|
bind_vars={"collection": c} |
|
|
) |
|
|
for _id in article_ids: |
|
|
self.arango_ids.append(_id) |
|
|
|
|
|
# A standard LLM for normal chat |
|
|
self.chatbot = LLM(messages=self.chat.chat_history2bot()) |
|
|
# A helper bot for generating queries or short prompts |
|
|
self.helperbot = LLM( |
|
|
temperature=0, |
|
|
model="small", |
|
|
max_length_answer=500, |
|
|
system_message=get_query_builder_system_message(), |
|
|
messages=self.chat.chat_history2bot(n_messages=4, remove_system=True), |
|
|
) |
|
|
# A specialized LLM picking which tool to use |
|
|
self.toolbot = LLM( |
|
|
temperature=0, |
|
|
system_message=""" |
|
|
You are an assistant bot helping an answering bot to answer a user's messages. |
|
|
Your task is to choose one or multiple tools that will help the answering bot to provide the user with the best possible answer. |
|
|
You should NEVER directly answer the user. You MUST choose a tool. |
|
|
""", |
|
|
chat=False, |
|
|
model="small", |
|
|
) |
|
|
|
|
|
# Load or register the passed-in tools |
|
|
if tools: |
|
|
self.tools = ToolRegistry.get_tools(tools=tools) |
|
|
else: |
|
|
self.tools = ToolRegistry.get_tools() |
|
|
|
|
|
# Store other kwargs |
|
|
for arg in kwargs: |
|
|
setattr(self, arg, kwargs[arg]) |
|
|
|
|
|
def get_chunks( |
|
|
self, |
|
|
user_input, |
|
|
collections=["sci_articles", "other_documents"], |
|
|
n_results=7, |
|
|
n_sources=4, |
|
|
filter=True, |
|
|
): |
|
|
# Basic version without Streamlit calls |
|
|
query = self.helperbot.generate( |
|
|
get_generate_vector_query_prompt(user_input, self.chat.role) |
|
|
).content.strip('"') |
|
|
|
|
|
combined_chunks = [] |
|
|
if collections: |
|
|
for collection in collections: |
|
|
where_filter = {"_id": {"$in": self.arango_ids}} if filter else {} |
|
|
chunks = self.get_chromadb().query( |
|
|
query=query, |
|
|
collection=collection, |
|
|
n_results=n_results, |
|
|
n_sources=n_sources, |
|
|
where=where_filter, |
|
|
max_retries=3, |
|
|
) |
|
|
for doc, meta, dist in zip( |
|
|
chunks["documents"][0], |
|
|
chunks["metadatas"][0], |
|
|
chunks["distances"][0], |
|
|
): |
|
|
combined_chunks.append( |
|
|
{"document": doc, "metadata": meta, "distance": dist} |
|
|
) |
|
|
combined_chunks.sort(key=lambda x: x["distance"]) |
|
|
|
|
|
# Keep the best chunks according to n_sources |
|
|
sources = set() |
|
|
closest_chunks = [] |
|
|
for chunk in combined_chunks: |
|
|
source_id = chunk["metadata"].get("_id", "no_id") |
|
|
if source_id not in sources: |
|
|
sources.add(source_id) |
|
|
closest_chunks.append(chunk) |
|
|
if len(sources) >= n_sources: |
|
|
break |
|
|
if len(closest_chunks) < n_results: |
|
|
remaining_chunks = [ |
|
|
c for c in combined_chunks if c not in closest_chunks |
|
|
] |
|
|
closest_chunks.extend(remaining_chunks[: n_results - len(closest_chunks)]) |
|
|
|
|
|
# Now fetch real metadata from Arango using the new API |
|
|
for chunk in closest_chunks: |
|
|
_id = chunk["metadata"].get("_id") |
|
|
if not _id: |
|
|
continue |
|
|
|
|
|
try: |
|
|
# Determine which database to use based on collection name |
|
|
if _id.startswith("sci_articles"): |
|
|
# Use base_arango for common documents |
|
|
arango_doc = self.base_arango.get_document(_id) |
|
|
else: |
|
|
# Use user_arango for user-specific documents |
|
|
arango_doc = self.user_arango.get_document(_id) |
|
|
|
|
|
if arango_doc: |
|
|
arango_metadata = arango_doc.get("metadata", {}) |
|
|
# Possibly merge notes |
|
|
if "user_notes" in arango_doc: |
|
|
arango_metadata["user_notes"] = arango_doc["user_notes"] |
|
|
chunk["metadata"] = arango_metadata |
|
|
except Exception as e: |
|
|
print_red(f"Error fetching document {_id}: {e}") |
|
|
|
|
|
# Group by article title |
|
|
grouped_chunks = {} |
|
|
article_number = 1 |
|
|
for chunk in closest_chunks: |
|
|
title = chunk["metadata"].get("title", "No title") |
|
|
chunk["article_number"] = article_number |
|
|
if title not in grouped_chunks: |
|
|
grouped_chunks[title] = { |
|
|
"article_number": article_number, |
|
|
"chunks": [], |
|
|
} |
|
|
article_number += 1 |
|
|
grouped_chunks[title]["chunks"].append(chunk) |
|
|
return grouped_chunks |
|
|
|
|
|
def answer_tool_call(self, response, user_input): |
|
|
bot_responses = [] |
|
|
# This method returns / stores responses (no Streamlit calls) |
|
|
if not response.get("tool_calls"): |
|
|
return "" |
|
|
|
|
|
for tool in response.get("tool_calls"): |
|
|
function_name = tool.function.get('name') |
|
|
arguments = tool.function.arguments |
|
|
arguments["query"] = user_input |
|
|
|
|
|
if hasattr(self, function_name): |
|
|
if function_name in [ |
|
|
"fetch_other_documents_tool", |
|
|
"fetch_science_articles_tool", |
|
|
"fetch_science_articles_and_other_documents_tool", |
|
|
]: |
|
|
chunks = getattr(self, function_name)(**arguments) |
|
|
bot_responses.append( |
|
|
self.generate_from_chunks(user_input, chunks).strip('"') |
|
|
) |
|
|
elif function_name == "fetch_notes_tool": |
|
|
notes = getattr(self, function_name)() |
|
|
bot_responses.append( |
|
|
self.generate_from_notes(user_input, notes).strip('"') |
|
|
) |
|
|
elif function_name == "conversational_response_tool": |
|
|
bot_responses.append( |
|
|
getattr(self, function_name)(user_input).strip('"') |
|
|
) |
|
|
return "\n\n".join(bot_responses) |
|
|
|
|
|
def process_user_input(self, user_input, content_attachment=None): |
|
|
# Add user message |
|
|
self.chat.add_message("user", user_input) |
|
|
|
|
|
if not content_attachment: |
|
|
prompt = get_tools_prompt(user_input) |
|
|
response = self.toolbot.generate(prompt, tools=self.tools, stream=False) |
|
|
if response.get("tool_calls"): |
|
|
bot_response = self.answer_tool_call(response, user_input) |
|
|
else: |
|
|
# Just respond directly |
|
|
bot_response = response.content.strip('"') |
|
|
else: |
|
|
# If there's an attachment, do something minimal |
|
|
bot_response = "Content attachment received (Base Bot)." |
|
|
|
|
|
# Add assistant message |
|
|
if self.chat.chat_history[-1]["role"] != "assistant": |
|
|
self.chat.add_message("assistant", bot_response) |
|
|
|
|
|
# Update in Arango |
|
|
self.chat.update_in_arango() |
|
|
return bot_response |
|
|
|
|
|
def generate_from_notes(self, user_input, notes): |
|
|
# No Streamlit calls |
|
|
notes_string = "" |
|
|
for note in notes: |
|
|
notes_string += f"\n# {note.get('title','No title')}\n{note.get('text','')}\n---\n" |
|
|
prompt = get_chat_prompt(user_input, content_string=notes_string, role=self.chat.role) |
|
|
return self.chatbot.generate(prompt, stream=True) |
|
|
|
|
|
def generate_from_chunks(self, user_input, chunks): |
|
|
# No Streamlit calls |
|
|
chunks_string = "" |
|
|
for title, group in chunks.items(): |
|
|
user_notes_string = "" |
|
|
if "user_notes" in group["chunks"][0]["metadata"]: |
|
|
notes = group["chunks"][0]["metadata"]["user_notes"] |
|
|
user_notes_string = f'\n\nUser notes:\n"""\n{notes}\n"""\n\n' |
|
|
docs = "\n(...)\n".join([c["document"] for c in group["chunks"]]) |
|
|
chunks_string += ( |
|
|
f"\n# {title}\n## Article #{group['article_number']}\n{user_notes_string}{docs}\n---\n" |
|
|
) |
|
|
prompt = get_chat_prompt(user_input, content_string=chunks_string, role=self.chat.role) |
|
|
return self.chatbot.generate(prompt, stream=True) |
|
|
|
|
|
def run(self): |
|
|
# Base Bot has no Streamlit run loop |
|
|
pass |
|
|
|
|
|
def get_notes(self): |
|
|
# Get project notes using the new API |
|
|
if self.project and hasattr(self.project, "name"): |
|
|
notes = self.user_arango.get_project_notes( |
|
|
project_name=self.project.name, |
|
|
username=self.username |
|
|
) |
|
|
return list(notes) |
|
|
return [] |
|
|
|
|
|
@ToolRegistry.register |
|
|
def fetch_science_articles_tool(self, query: str, n_documents: int): |
|
|
""" |
|
|
"Fetches information from scientific articles. Use this tool when the user is looking for information from scientific articles." |
|
|
|
|
|
Parameters: |
|
|
query (str): The search query to find relevant scientific articles. |
|
|
n_documents (int): How many documents to fetch. A complex query may require more documents. Min: 3, Max: 10. |
|
|
|
|
|
Returns: |
|
|
list: A list of chunks containing information from the fetched scientific articles. |
|
|
""" |
|
|
print_purple('Query:', query) |
|
|
|
|
|
n_documents = int(n_documents) |
|
|
if n_documents < 3: |
|
|
n_documents = 3 |
|
|
elif n_documents > 10: |
|
|
n_documents = 10 |
|
|
return self.get_chunks( |
|
|
query, collections=["sci_articles"], n_results=n_documents |
|
|
) |
|
|
|
|
|
@ToolRegistry.register |
|
|
def fetch_other_documents_tool(self, query: str, n_documents: int): |
|
|
""" |
|
|
Fetches information from other documents based on the user's query. |
|
|
|
|
|
This method retrieves information from various types of documents such as reports, news articles, and other texts. It should be used only when it is clear that the user is not seeking scientific articles. |
|
|
|
|
|
Args: |
|
|
query (str): The search query provided by the user. |
|
|
n_documents (int): How many documents to fetch. A complex query may require more documents. Min: 2, Max: 10. |
|
|
|
|
|
Returns: |
|
|
list: A list of document chunks that match the query. |
|
|
""" |
|
|
assert isinstance(self, Bot), "The first argument must be a Bot object." |
|
|
n_documents = int(n_documents) |
|
|
if n_documents < 2: |
|
|
n_documents = 2 |
|
|
elif n_documents > 10: |
|
|
n_documents = 10 |
|
|
return self.get_chunks( |
|
|
query, |
|
|
collections=[f"{self.username}__other_documents"], |
|
|
n_results=n_documents, |
|
|
) |
|
|
|
|
|
@ToolRegistry.register |
|
|
def fetch_science_articles_and_other_documents_tool( |
|
|
self, query: str, n_documents: int |
|
|
): |
|
|
""" |
|
|
Fetches information from both scientific articles and other documents. |
|
|
|
|
|
This method is often used when the user hasn't specified what kind of sources they are interested in. |
|
|
|
|
|
Args: |
|
|
query (str): The search query to fetch information for. |
|
|
n_documents (int): How many documents to fetch. A complex query may require more documents. Min: 3, Max: 10. |
|
|
|
|
|
Returns: |
|
|
list: A list of document chunks that match the search query. |
|
|
""" |
|
|
assert isinstance(self, Bot), "The first argument must be a Bot object." |
|
|
n_documents = int(n_documents) |
|
|
if n_documents < 3: |
|
|
n_documents = 3 |
|
|
elif n_documents > 10: |
|
|
n_documents = 10 |
|
|
return self.get_chunks( |
|
|
query, |
|
|
collections=["sci_articles", f"{self.username}__other_documents"], |
|
|
n_results=n_documents, |
|
|
) |
|
|
|
|
|
@ToolRegistry.register |
|
|
def fetch_notes_tool(bot): |
|
|
""" |
|
|
Fetches information from the project notes when you as an editor need context from the project notes to understand other information. ONLY use this together with other tools! No arguments needed. |
|
|
|
|
|
Returns: |
|
|
list: A list of notes. |
|
|
""" |
|
|
assert isinstance(bot, Bot), "The first argument must be a Bot object." |
|
|
return bot.get_notes() |
|
|
|
|
|
@ToolRegistry.register |
|
|
def conversational_response_tool(self, query: str): |
|
|
""" |
|
|
Generate a conversational response to a user's query. |
|
|
|
|
|
This method is designed to provide a short and conversational response |
|
|
without fetching additional data. It should be used only when it is clear |
|
|
that the user is engaging in small talk (like saying 'hi') and not seeking detailed information. |
|
|
|
|
|
Args: |
|
|
query (str): The user's message to which the bot should respond. |
|
|
|
|
|
Returns: |
|
|
str: The generated conversational response. |
|
|
""" |
|
|
query = f""" |
|
|
User message: "{query}". |
|
|
Make your answer short and conversational. |
|
|
This is perhaps not a conversation about a journalistic project, so try not to be too informative. |
|
|
Don't answer with anything you're not sure of! |
|
|
""" |
|
|
|
|
|
result = ( |
|
|
self.chatbot.generate(query, stream=True) |
|
|
if self.chatbot |
|
|
else self.llm.generate(query, stream=True) |
|
|
) |
|
|
return result
|
|
|
|
