wehub-resource-sync
444 lines
17 KiB
Python
444 lines
17 KiB
Python
import nest_asyncio
|
|
nest_asyncio.apply()
|
|
|
|
import os
|
|
import asyncio
|
|
import streamlit as st
|
|
import base64
|
|
import gc
|
|
import tempfile
|
|
import uuid
|
|
import time
|
|
import io
|
|
import re
|
|
from contextlib import redirect_stdout
|
|
from pathlib import Path
|
|
|
|
from src.embeddings.embed_data import EmbedData
|
|
from src.indexing.milvus_vdb import MilvusVDB
|
|
from src.retrieval.retriever_rerank import Retriever
|
|
from src.generation.rag import RAG
|
|
from src.workflows.agent_workflow import ParalegalAgentWorkflow
|
|
from pypdf import PdfReader
|
|
from dotenv import load_dotenv
|
|
from config.settings import settings
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Set up page configuration
|
|
st.set_page_config(page_title="Paralegal AI Assistant", layout="wide")
|
|
|
|
# Initialize session state variables
|
|
if "id" not in st.session_state:
|
|
st.session_state.id = str(uuid.uuid4())[:8]
|
|
st.session_state.file_cache = {}
|
|
|
|
if "workflow" not in st.session_state:
|
|
st.session_state.workflow = None
|
|
|
|
if "messages" not in st.session_state:
|
|
st.session_state.messages = []
|
|
|
|
if "workflow_logs" not in st.session_state:
|
|
st.session_state.workflow_logs = []
|
|
|
|
if "vector_db" not in st.session_state:
|
|
st.session_state.vector_db = None
|
|
|
|
session_id = st.session_state.id
|
|
|
|
def reset_chat():
|
|
"""Reset chat history and clear memory."""
|
|
st.session_state.messages = []
|
|
st.session_state.workflow_logs = []
|
|
gc.collect()
|
|
|
|
def display_pdf(file):
|
|
"""Display PDF preview in sidebar."""
|
|
st.markdown("### PDF Preview")
|
|
base64_pdf = base64.b64encode(file.read()).decode("utf-8")
|
|
|
|
pdf_display = f"""<iframe src="data:application/pdf;base64,{base64_pdf}" width="400" height="100%" type="application/pdf"
|
|
style="height:100vh; width:100%"
|
|
>
|
|
</iframe>"""
|
|
|
|
st.markdown(pdf_display, unsafe_allow_html=True)
|
|
|
|
def render_logs(log_text: str):
|
|
"""Render logs with ANSI colors and emojis nicely in Streamlit"""
|
|
from ansi2html import Ansi2HTMLConverter
|
|
conv = Ansi2HTMLConverter(inline=True)
|
|
html_body = conv.convert(log_text, full=False)
|
|
|
|
st.markdown(
|
|
f"""
|
|
<div style="font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', 'Courier New', monospace; white-space: pre-wrap; line-height: 1.45; font-size: 13px;">
|
|
{html_body}
|
|
</div>
|
|
""",
|
|
unsafe_allow_html=True,
|
|
)
|
|
|
|
def load_and_split_pdf(file_path: str, chunk_size: int = 512, chunk_overlap: int = 50):
|
|
try:
|
|
reader = PdfReader(file_path)
|
|
full_text_parts = []
|
|
for page in reader.pages:
|
|
text = page.extract_text() or ""
|
|
if text:
|
|
full_text_parts.append(text)
|
|
full_text = "\n".join(full_text_parts)
|
|
|
|
words = full_text.split()
|
|
chunks = []
|
|
i = 0
|
|
step = max(1, chunk_size - chunk_overlap)
|
|
while i < len(words):
|
|
segment = words[i : i + chunk_size]
|
|
chunks.append(" ".join(segment))
|
|
i += step
|
|
return [c for c in chunks if c.strip()]
|
|
except Exception as e:
|
|
st.error(f"Error loading PDF: {e}")
|
|
return []
|
|
|
|
def initialize_workflow(file_path: str):
|
|
with st.spinner("🔄 Loading document and setting up the workflow..."):
|
|
try:
|
|
# Step 1: Load and split document
|
|
st.info("📄 Loading and processing PDF...")
|
|
text_chunks = load_and_split_pdf(file_path)
|
|
|
|
if not text_chunks:
|
|
st.error("No text chunks extracted from PDF")
|
|
return None
|
|
|
|
st.success(f"✅ Created {len(text_chunks)} text chunks")
|
|
|
|
# Step 2: Create embeddings
|
|
st.info("🧠 Generating embeddings...")
|
|
embed_data = EmbedData(
|
|
embed_model_name=settings.embedding_model,
|
|
batch_size=settings.batch_size
|
|
)
|
|
embed_data.embed(text_chunks)
|
|
st.success("✅ Embeddings generated with binary quantization")
|
|
|
|
# Step 3: Setup vector database
|
|
st.info("🗄️ Setting up Milvus vector database...")
|
|
collection_name = f"{settings.collection_name}_{session_id}"
|
|
|
|
vector_db = MilvusVDB(
|
|
collection_name=collection_name,
|
|
vector_dim=settings.vector_dim,
|
|
batch_size=settings.batch_size,
|
|
db_file=f"{settings.milvus_db_path}_{session_id}.db"
|
|
)
|
|
|
|
vector_db.initialize_client()
|
|
vector_db.create_collection()
|
|
vector_db.ingest_data(embed_data)
|
|
|
|
# Store in session state for cleanup
|
|
st.session_state.vector_db = vector_db
|
|
st.success("✅ Vector database setup completed")
|
|
|
|
# Step 4: Setup retrieval
|
|
st.info("🔍 Setting up retrieval system...")
|
|
retriever = Retriever(
|
|
vector_db=vector_db,
|
|
embed_data=embed_data,
|
|
top_k=settings.top_k
|
|
)
|
|
st.success("✅ Retrieval system ready")
|
|
|
|
# Step 5: Setup RAG system
|
|
st.info("🤖 Setting up RAG system...")
|
|
rag_system = RAG(
|
|
retriever=retriever,
|
|
llm_model=settings.llm_model,
|
|
temperature=settings.temperature,
|
|
max_tokens=settings.max_tokens
|
|
)
|
|
st.success("✅ RAG system initialized")
|
|
|
|
# Step 6: Setup workflow
|
|
st.info("⚙️ Setting up agentic workflow...")
|
|
workflow = ParalegalAgentWorkflow(
|
|
retriever=retriever,
|
|
rag_system=rag_system,
|
|
firecrawl_api_key=settings.firecrawl_api_key or os.getenv("FIRECRAWL_API_KEY"),
|
|
openai_api_key=settings.openai_api_key or os.getenv("OPENAI_API_KEY")
|
|
)
|
|
|
|
st.success("🎉 Workflow setup completed!")
|
|
return workflow
|
|
|
|
except Exception as e:
|
|
st.error(f"Error initializing workflow: {e}")
|
|
return None
|
|
|
|
async def run_workflow(query: str):
|
|
f = io.StringIO()
|
|
with redirect_stdout(f):
|
|
result = await st.session_state.workflow.run_workflow(query)
|
|
|
|
# Get aptured logs and store them
|
|
logs = f.getvalue()
|
|
if logs:
|
|
st.session_state.workflow_logs.append(logs)
|
|
|
|
return result
|
|
|
|
def cleanup_resources():
|
|
"""Cleanup vector database and other resources."""
|
|
if st.session_state.vector_db:
|
|
try:
|
|
st.session_state.vector_db.close()
|
|
except:
|
|
pass
|
|
st.session_state.vector_db = None
|
|
|
|
# Sidebar for configuration and document upload
|
|
with st.sidebar:
|
|
st.header("🔧 Configuration")
|
|
|
|
# st.subheader("API Keys")
|
|
# openai_key = st.text_input("OpenAI API Key", type="password", value=os.getenv("OPENAI_API_KEY", ""))
|
|
ollama_model = st.text_input("Ollama Model", value="gpt-oss:20b")
|
|
firecrawl_key = st.text_input("Firecrawl API Key", type="password", value=os.getenv("FIRECRAWL_API_KEY", ""))
|
|
|
|
# if openai_key:
|
|
# os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
|
|
# st.success("✅ OpenAI API Key set!")
|
|
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
|
|
if firecrawl_key:
|
|
os.environ["FIRECRAWL_API_KEY"] = firecrawl_key
|
|
st.success("✅ Firecrawl API Key set!")
|
|
|
|
st.markdown("---")
|
|
|
|
# Document upload section
|
|
st.header("📄 Upload Document")
|
|
st.markdown("Upload a PDF document to get started")
|
|
|
|
uploaded_file = st.file_uploader("Choose your PDF file", type="pdf")
|
|
|
|
if uploaded_file:
|
|
try:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
file_path = os.path.join(temp_dir, uploaded_file.name)
|
|
|
|
with open(file_path, "wb") as f:
|
|
f.write(uploaded_file.getvalue())
|
|
|
|
file_key = f"{session_id}-{uploaded_file.name}"
|
|
|
|
if file_key not in st.session_state.get('file_cache', {}):
|
|
# Initialize workflow with the uploaded document
|
|
workflow = initialize_workflow(file_path)
|
|
if workflow:
|
|
st.session_state.workflow = workflow
|
|
st.session_state.file_cache[file_key] = workflow
|
|
st.balloons()
|
|
else:
|
|
st.session_state.workflow = st.session_state.file_cache[file_key]
|
|
|
|
if st.session_state.workflow:
|
|
st.success("🎉 Ready to Chat!")
|
|
display_pdf(uploaded_file)
|
|
|
|
except Exception as e:
|
|
st.error(f"An error occurred: {e}")
|
|
|
|
# elif uploaded_file and not openai_key:
|
|
# st.warning("⚠️ Please enter your OpenAI API key first!")
|
|
elif uploaded_file:
|
|
st.info("📁 Please upload a PDF to continue")
|
|
|
|
# Cleanup button
|
|
st.markdown("---")
|
|
if st.button("🗑️ Clean Up Resources"):
|
|
cleanup_resources()
|
|
st.success("Resources cleaned up!")
|
|
|
|
# Main chat interface
|
|
col1, col2 = st.columns([6, 1])
|
|
|
|
with col1:
|
|
st.markdown('''
|
|
<h1 style='color: #2E86AB; margin-bottom: 10px;'>
|
|
⚖️ Paralegal AI assistant
|
|
</h1>
|
|
<div style="display: flex; align-items: center; gap: 8px; margin-bottom: 20px;">
|
|
<span style='color: #A23B72; font-size: 16px;'>Powered by</span>
|
|
<div style="display: flex; align-items: center; gap: 20px;">
|
|
<a href="#" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://images.seeklogo.com/logo-png/61/2/crew-ai-logo-png_seeklogo-619843.png"
|
|
alt="CrewAI" style="height: 100px;">
|
|
</a>
|
|
<a href="#" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://milvus.io/images/layout/milvus-logo.svg"
|
|
alt="Milvus" style="height: 32px;">
|
|
</a>
|
|
<a href="#" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://i.ibb.co/VcsfddTr/logo-dark.png"
|
|
alt="Firecrawl" style="height: 45px;">
|
|
</a>
|
|
<a href="#" style="display: inline-block; vertical-align: middle;">
|
|
<img src="https://i.ibb.co/wt57zN1/ollama.png"
|
|
alt="Ollama" style="height: 48px;">
|
|
</a>
|
|
</div>
|
|
</div>
|
|
''', unsafe_allow_html=True)
|
|
|
|
with col2:
|
|
if st.button("Clear Chat ↺", on_click=reset_chat):
|
|
st.rerun()
|
|
|
|
# System info
|
|
if st.session_state.workflow:
|
|
st.success("🟢 System Ready - Workflow initialized successfully!")
|
|
else:
|
|
st.info("🔵 Upload a PDF document to get started")
|
|
|
|
# Display chat messages from history
|
|
for i, message in enumerate(st.session_state.messages):
|
|
with st.chat_message(message["role"]):
|
|
st.markdown(message["content"])
|
|
|
|
# # Display workflow logs for user messages
|
|
# if (message["role"] == "user" and
|
|
# "log_index" in message and
|
|
# message["log_index"] < len(st.session_state.workflow_logs)):
|
|
|
|
# with st.expander("🔍 View Workflow Execution Details", expanded=False):
|
|
# logs = st.session_state.workflow_logs[message["log_index"]]
|
|
# render_logs(logs)
|
|
|
|
# Accept user input
|
|
if prompt := st.chat_input("Ask a question about your document..."):
|
|
if not st.session_state.workflow:
|
|
st.error("⚠️ Please upload a document first to initialize the workflow.")
|
|
st.stop()
|
|
|
|
if not os.getenv("OPENAI_API_KEY"):
|
|
st.error("⚠️ Please set your OpenAI API key in the sidebar.")
|
|
st.stop()
|
|
|
|
# Add user message to chat history
|
|
log_index = len(st.session_state.workflow_logs)
|
|
st.session_state.messages.append({
|
|
"role": "user",
|
|
"content": prompt,
|
|
"log_index": log_index
|
|
})
|
|
|
|
# Display user message
|
|
with st.chat_message("user"):
|
|
st.markdown(prompt)
|
|
|
|
# Run the workflow and get response
|
|
with st.chat_message("assistant"):
|
|
message_placeholder = st.empty()
|
|
|
|
try:
|
|
with st.spinner("🔄 Processing your query..."):
|
|
# Measure end-to-end workflow time
|
|
workflow_start = time.perf_counter()
|
|
result = asyncio.run(run_workflow(prompt))
|
|
workflow_end = time.perf_counter()
|
|
workflow_time = workflow_end - workflow_start
|
|
|
|
# # Display workflow logs
|
|
# if log_index < len(st.session_state.workflow_logs):
|
|
# with st.expander("🔍 View Workflow Execution Details", expanded=False):
|
|
# render_logs(st.session_state.workflow_logs[log_index])
|
|
|
|
# Get the final answer
|
|
if isinstance(result, dict) and "answer" in result:
|
|
full_response = result["answer"]
|
|
|
|
# Show additional info about the workflow
|
|
if result.get("web_search_used", False):
|
|
st.info("🌐 This response includes information from web search")
|
|
# if 'workflow_time' in locals():
|
|
# st.caption(f"🕒 Completion time: {workflow_time:.2f} s")
|
|
else:
|
|
st.info("📚 This response is based on your document")
|
|
try:
|
|
retriever = getattr(st.session_state.workflow, "retriever", None)
|
|
if retriever:
|
|
retrieve_start = time.perf_counter()
|
|
retriever.search(prompt)
|
|
retrieve_end = time.perf_counter()
|
|
retrieval_time = retrieve_end - retrieve_start
|
|
|
|
citations = retriever.get_citations(prompt, top_k=settings.top_k, snippet_chars=300)
|
|
|
|
if citations:
|
|
with st.expander("📎 Citations (top matches)"):
|
|
for c in citations:
|
|
score = c.get("score")
|
|
try:
|
|
score_str = f"{float(score):.3f}"
|
|
except Exception:
|
|
score_str = str(score)
|
|
st.markdown(
|
|
f"[{c['rank']}] score={score_str} id={c.get('node_id')}"
|
|
)
|
|
if c.get("snippet"):
|
|
st.code(c["snippet"], language="text")
|
|
except Exception as e:
|
|
st.warning(f"Could not fetch citations: {e}")
|
|
|
|
# Show timing caption
|
|
times = []
|
|
if retrieval_time is not None:
|
|
times.append(f"🕒 Retrieval time: {retrieval_time:.2f} s")
|
|
# if 'workflow_time' in locals():
|
|
# times.append(f"🕒 Completion time: {workflow_time:.2f} s")
|
|
if times:
|
|
st.caption(" • ".join(times))
|
|
|
|
else:
|
|
full_response = str(result)
|
|
|
|
# Stream the response word by word
|
|
streamed_response = ""
|
|
words = full_response.split()
|
|
|
|
for i, word in enumerate(words):
|
|
streamed_response += word + " "
|
|
message_placeholder.markdown(streamed_response + "▌")
|
|
|
|
if i < len(words) - 1:
|
|
time.sleep(0.05)
|
|
|
|
# Display final response
|
|
message_placeholder.markdown(full_response)
|
|
|
|
except Exception as e:
|
|
error_msg = f"❌ Error processing your question: {str(e)}"
|
|
st.error(error_msg)
|
|
full_response = "I apologize, but I encountered an error while processing your question. Please try again."
|
|
message_placeholder.markdown(full_response)
|
|
|
|
# Add assistant response to chat history
|
|
st.session_state.messages.append({
|
|
"role": "assistant",
|
|
"content": full_response
|
|
})
|
|
|
|
# Footer
|
|
st.markdown("---")
|
|
st.markdown(
|
|
"<p style='text-align: center; color: #666; font-size: 12px;'>"
|
|
"Paralegal AI assistant • Built with Streamlit, CrewAI, Milvus, Firecrawl, and Ollama"
|
|
"</p>",
|
|
unsafe_allow_html=True
|
|
)
|