chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,443 @@
|
||||
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
|
||||
)
|
||||
Reference in New Issue
Block a user