wehub-resource-sync
187 lines
6.7 KiB
Python
187 lines
6.7 KiB
Python
import streamlit as st
|
|
import os
|
|
import tempfile
|
|
import gc
|
|
import base64
|
|
import time
|
|
|
|
from crewai import Agent, Crew, Process, Task
|
|
from crewai_tools import SerperDevTool
|
|
from src.agentic_rag.tools.custom_tool import DocumentSearchTool
|
|
|
|
# ===========================
|
|
# Define Agents & Tasks
|
|
# ===========================
|
|
def create_agents_and_tasks(pdf_tool):
|
|
"""Creates a Crew with the given PDF tool (if any) and a web search tool."""
|
|
web_search_tool = SerperDevTool()
|
|
|
|
retriever_agent = Agent(
|
|
role="Retrieve relevant information to answer the user query: {query}",
|
|
goal=(
|
|
"Retrieve the most relevant information from the available sources "
|
|
"for the user query: {query}. Always try to use the PDF search tool first. "
|
|
"If you are not able to retrieve the information from the PDF search tool, "
|
|
"then try to use the web search tool."
|
|
),
|
|
backstory=(
|
|
"You're a meticulous analyst with a keen eye for detail. "
|
|
"You're known for your ability to understand user queries: {query} "
|
|
"and retrieve knowledge from the most suitable knowledge base."
|
|
),
|
|
verbose=True,
|
|
tools=[t for t in [pdf_tool, web_search_tool] if t],
|
|
)
|
|
|
|
response_synthesizer_agent = Agent(
|
|
role="Response synthesizer agent for the user query: {query}",
|
|
goal=(
|
|
"Synthesize the retrieved information into a concise and coherent response "
|
|
"based on the user query: {query}. If you are not able to retrieve the "
|
|
'information then respond with "I\'m sorry, I couldn\'t find the information '
|
|
'you\'re looking for."'
|
|
),
|
|
backstory=(
|
|
"You're a skilled communicator with a knack for turning "
|
|
"complex information into clear and concise responses."
|
|
),
|
|
verbose=True
|
|
)
|
|
|
|
retrieval_task = Task(
|
|
description=(
|
|
"Retrieve the most relevant information from the available "
|
|
"sources for the user query: {query}"
|
|
),
|
|
expected_output=(
|
|
"The most relevant information in the form of text as retrieved "
|
|
"from the sources."
|
|
),
|
|
agent=retriever_agent
|
|
)
|
|
|
|
response_task = Task(
|
|
description="Synthesize the final response for the user query: {query}",
|
|
expected_output=(
|
|
"A concise and coherent response based on the retrieved information "
|
|
"from the right source for the user query: {query}. If you are not "
|
|
"able to retrieve the information, then respond with: "
|
|
'"I\'m sorry, I couldn\'t find the information you\'re looking for."'
|
|
),
|
|
agent=response_synthesizer_agent
|
|
)
|
|
|
|
crew = Crew(
|
|
agents=[retriever_agent, response_synthesizer_agent],
|
|
tasks=[retrieval_task, response_task],
|
|
process=Process.sequential, # or Process.hierarchical
|
|
verbose=True
|
|
)
|
|
return crew
|
|
|
|
# ===========================
|
|
# Streamlit Setup
|
|
# ===========================
|
|
if "messages" not in st.session_state:
|
|
st.session_state.messages = [] # Chat history
|
|
|
|
if "pdf_tool" not in st.session_state:
|
|
st.session_state.pdf_tool = None # Store the DocumentSearchTool
|
|
|
|
if "crew" not in st.session_state:
|
|
st.session_state.crew = None # Store the Crew object
|
|
|
|
def reset_chat():
|
|
st.session_state.messages = []
|
|
gc.collect()
|
|
|
|
def display_pdf(file_bytes: bytes, file_name: str):
|
|
"""Displays the uploaded PDF in an iframe."""
|
|
base64_pdf = base64.b64encode(file_bytes).decode("utf-8")
|
|
pdf_display = f"""
|
|
<iframe
|
|
src="data:application/pdf;base64,{base64_pdf}"
|
|
width="100%"
|
|
height="600px"
|
|
type="application/pdf"
|
|
>
|
|
</iframe>
|
|
"""
|
|
st.markdown(f"### Preview of {file_name}")
|
|
st.markdown(pdf_display, unsafe_allow_html=True)
|
|
|
|
# ===========================
|
|
# Sidebar
|
|
# ===========================
|
|
with st.sidebar:
|
|
st.header("Add Your PDF Document")
|
|
uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"])
|
|
|
|
if uploaded_file is not None:
|
|
# If there's a new file and we haven't set pdf_tool yet...
|
|
if st.session_state.pdf_tool is None:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
temp_file_path = os.path.join(temp_dir, uploaded_file.name)
|
|
with open(temp_file_path, "wb") as f:
|
|
f.write(uploaded_file.getvalue())
|
|
|
|
with st.spinner("Indexing PDF... Please wait..."):
|
|
st.session_state.pdf_tool = DocumentSearchTool(file_path=temp_file_path)
|
|
|
|
st.success("PDF indexed! Ready to chat.")
|
|
|
|
# Optionally display the PDF in the sidebar
|
|
display_pdf(uploaded_file.getvalue(), uploaded_file.name)
|
|
|
|
st.button("Clear Chat", on_click=reset_chat)
|
|
|
|
# ===========================
|
|
# Main Chat Interface
|
|
# ===========================
|
|
st.markdown("""
|
|
# Agentic RAG powered by <img src="data:image/png;base64,{}" width="120" style="vertical-align: -3px;">
|
|
""".format(base64.b64encode(open("assets/crewai.png", "rb").read()).decode()), unsafe_allow_html=True)
|
|
|
|
# Render existing conversation
|
|
for message in st.session_state.messages:
|
|
with st.chat_message(message["role"]):
|
|
st.markdown(message["content"])
|
|
|
|
# Chat input
|
|
prompt = st.chat_input("Ask a question about your PDF...")
|
|
|
|
if prompt:
|
|
# 1. Show user message immediately
|
|
st.session_state.messages.append({"role": "user", "content": prompt})
|
|
with st.chat_message("user"):
|
|
st.markdown(prompt)
|
|
|
|
# 2. Build or reuse the Crew (only once after PDF is loaded)
|
|
if st.session_state.crew is None:
|
|
st.session_state.crew = create_agents_and_tasks(st.session_state.pdf_tool)
|
|
|
|
# 3. Get the response
|
|
with st.chat_message("assistant"):
|
|
message_placeholder = st.empty()
|
|
full_response = ""
|
|
|
|
# Get the complete response first
|
|
with st.spinner("Thinking..."):
|
|
inputs = {"query": prompt}
|
|
result = st.session_state.crew.kickoff(inputs=inputs).raw
|
|
|
|
# Split by lines first to preserve code blocks and other markdown
|
|
lines = result.split('\n')
|
|
for i, line in enumerate(lines):
|
|
full_response += line
|
|
if i < len(lines) - 1: # Don't add newline to the last line
|
|
full_response += '\n'
|
|
message_placeholder.markdown(full_response + "▌")
|
|
time.sleep(0.15) # Adjust the speed as needed
|
|
|
|
# Show the final response without the cursor
|
|
message_placeholder.markdown(full_response)
|
|
|
|
# 4. Save assistant's message to session
|
|
st.session_state.messages.append({"role": "assistant", "content": result})
|