chore: import upstream snapshot with attribution

paralegal-agent-crew/.env.example

@@ -0,0 +1,2 @@
+OPENAI_API_KEY=<YOUR_OPENAI_API_KEY>
+FIRECRAWL_API_KEY=<YOUR_FIRECRAWL_API_KEY>

paralegal-agent-crew/.gitignore

@@ -0,0 +1,61 @@
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+*.pyd
+*.so
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Build / packaging
+build/
+dist/
+*.egg-info/
+.eggs/
+.wheels/
+
+# IDE / Editor
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# OS files
+.DS_Store
+Thumbs.db
+
+# Logs
+*.log
+logs/
+
+# Caches
+.cache/
+.pytest_cache/
+
+# Data and model caches
+cache/
+cache/**
+
+# Local databases and artifacts
+data/*.db
+data/*.db.lock
+*.db
+*.sqlite
+*.sqlite3
+
+# Streamlit
+.streamlit/
+
+# Env files / secrets
+.env
+*.env
+
+# uv / pip tools
+uv.lock
+pip-wheel-metadata/
+
+# Firecrawl temp
+.firecrawl/

paralegal-agent-crew/README.md

@@ -0,0 +1,93 @@
+# Paralegal AI Assistant
+
+⚖️ An intelligent paralegal AI assistant that analyzes PDF documents and provides comprehensive answers through advanced RAG (Retrieval-Augmented Generation) with web search fallback capabilities.
+
+## Setup Instructions
+
+### Prerequisites
+- Python 3.13+
+- OpenAI API key
+- Firecrawl API key (optional)
+- Docker (for self-hosted Milvus)
+
+### Installation
+
+1. **Clone and navigate to the project:**
+   ```bash
+   git clone <repository-url>
+   cd paralegal-agent-crew
+   ```
+
+2. **Set up environment variables:**
+   ```bash
+   cp .env.example .env
+   ```
+   Edit `.env` and add your API keys:
+   ```env
+   OPENAI_API_KEY=your_openai_api_key_here
+   FIRECRAWL_API_KEY=your_firecrawl_api_key_here
+   ```
+
+3. **Install dependencies with UV:**
+   ```bash
+   uv sync
+   ```
+
+### Option A: Using Local Milvus (Default)
+The project works with embedded Milvus out of the box. Simply run:
+
+```bash
+uv run streamlit run app.py
+```
+
+### Option B: Self-Hosted Milvus with Docker
+
+4. **Set up Milvus vector database:**
+   
+   **Quick Setup (Recommended):**
+   ```bash
+   # Download and run Milvus installation script
+   curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh -o standalone_embed.sh
+   bash standalone_embed.sh start
+   ```
+   
+   **Alternative - Docker Compose:**
+   ```bash
+   # Download docker-compose file
+   wget https://github.com/milvus-io/milvus/releases/download/v2.0.2/milvus-standalone-docker-compose.yml -O docker-compose.yml
+   
+   # Start Milvus
+   docker-compose up -d
+   ```
+
+5. **Update configuration for external Milvus:**
+   Modify `config/settings.py` to point to your Milvus instance:
+   ```python
+   milvus_host: str = "localhost"
+   milvus_port: int = 19530
+   ```
+
+6. **Run the application:**
+   ```bash
+   uv run streamlit run app.py
+   ```
+
+7. **Open your browser and go to `http://localhost:8501`**
+
+### Milvus Management
+- **Milvus WebUI**: Access at `http://127.0.0.1:9091/webui/`
+- **Stop Milvus**: `bash standalone_embed.sh stop`
+- **Delete Milvus**: `bash standalone_embed.sh delete`
+
+## About the Project
+
+This paralegal AI assistant combines multiple technologies to provide intelligent document analysis:
+
+- **Document Processing**: Extracts and chunks PDF documents for analysis
+- **Vector Database**: Uses Milvus with binary quantization for efficient similarity search
+- **Embeddings**: BGE-large-en-v1.5 model for high-quality text representations
+- **Intelligent Routing**: Automatically evaluates response quality and triggers web search when needed
+- **Web Search Integration**: Firecrawl integration for additional context from the web
+- **Workflow Management**: CrewAI-powered agentic workflows for complex query handling
+
+The system provides an interactive Streamlit interface where users can upload PDF documents, ask questions, and receive comprehensive answers with citations and sources. It automatically determines when to use document knowledge versus web search to provide the most accurate and complete responses.

paralegal-agent-crew/app.py

@@ -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
+)

paralegal-agent-crew/config/settings.py

@@ -0,0 +1,44 @@
+import os
+from pathlib import Path
+from pydantic_settings import BaseSettings
+
+class Settings(BaseSettings):
+    # API Keys
+    openai_api_key: str
+    firecrawl_api_key: str
+    
+    # Model Configuration
+    embedding_model: str = "BAAI/bge-large-en-v1.5"
+    llm_model: str = "gpt-3.5-turbo"
+    vector_dim: int = 1024
+    
+    # Retrieval Configuration
+    top_k: int = 3
+    batch_size: int = 512
+    rerank_top_k: int = 3
+    
+    # Database Configuration
+    milvus_db_path: str = "./data/milvus_binary.db"
+    collection_name: str = "paralegal_agent"
+    
+    # Data Configuration
+    docs_path: str = "./data/raft.pdf"
+
+    # Cache Configuration
+    hf_cache_dir: str = "./cache/hf_cache"
+    
+    # LLM settings
+    temperature: float = 0.6
+    max_tokens: int = 1000
+    
+    class Config:
+        env_file = ".env"
+        case_sensitive = False
+    
+    def __post_init__(self):
+        # Create necessary directories
+        Path(self.milvus_db_path).parent.mkdir(parents=True, exist_ok=True)
+        Path(self.hf_cache_dir).mkdir(parents=True, exist_ok=True)
+
+# Global settings instance
+settings = Settings()

paralegal-agent-crew/examples/test.py

@@ -0,0 +1,261 @@
+import asyncio
+import os
+import sys
+from pathlib import Path
+from loguru import logger
+import warnings
+warnings.filterwarnings("ignore")
+
+sys.path.append(str(Path(__file__).parent.parent))
+
+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 config.settings import settings
+
+def get_citations(retriever, query, top_k=3, snippet_chars=300):
+    """Return retrieval results as simple citation dicts."""
+    results = retriever.search_with_scores(query, top_k=top_k)
+    citations = []
+    for rank, item in enumerate(results, start=1):
+        context = (item.get("context") or "").strip()
+        snippet = (context[:snippet_chars] + ("…" if len(context) > snippet_chars else "")) if context else ""
+        citations.append({
+            "rank": rank,
+            "node_id": item.get("node_id"),
+            "score": item.get("score"),
+            "snippet": snippet,
+            "metadata": item.get("metadata") or {},
+        })
+    return citations
+
+def print_citations(citations):
+    if not citations:
+        print("\n(No citations available)")
+        return
+    print("\nCITATIONS (top matches)")
+    print("-" * 60)
+    for c in citations:
+        score_str = f"{float(c.get("score")):.3f}"
+        node_id = c.get("node_id")
+        snippet = c.get("snippet") or ""
+        print(f"[{c['rank']}] score={score_str} id={node_id}")
+        if snippet:
+            print(f"    \u201c{snippet}\u201d")
+    print("-" * 60)
+
+async def main():
+    logger.info("Starting Enhanced RAG Pipeline Demo")
+    
+    required_env_vars = ["OPENAI_API_KEY"]
+    for var in required_env_vars:
+        if not os.getenv(var):
+            logger.error(f"Missing required environment variable: {var}")
+            return
+    
+    try:
+        # Step 1: Load and process document
+        logger.info("Step 1: Loading document...")
+        
+        docs_path = settings.docs_path
+        if not docs_path or not Path(docs_path).exists():
+            logger.error("Invalid PDF path provided")
+            return
+        
+        # Load and split documents
+        reader = PdfReader(docs_path)
+        pages_text = []
+        for page in reader.pages:
+            pages_text.append(page.extract_text() or "")
+        full_text = "\n".join(pages_text)
+        words = full_text.split()
+        text_chunks = []
+        chunk_size = 512
+        overlap = 50
+        step = max(1, chunk_size - overlap)
+        i = 0
+        while i < len(words):
+            segment = words[i : i + chunk_size]
+            text_chunks.append(" ".join(segment))
+            i += step
+        logger.info(f"Created {len(text_chunks)} text chunks")
+        
+        # Step 2: Create embeddings
+        logger.info("Step 2: Creating embeddings...")
+        
+        embed_data = EmbedData(
+            embed_model_name=settings.embedding_model,
+            batch_size=settings.batch_size
+        )
+        
+        # Generate embeddings with binary quantization
+        embed_data.embed(text_chunks)
+        logger.info("Embeddings created successfully")
+        
+        # Step 3: Setup vector database
+        logger.info("Step 3: Setting up vector database...")
+        
+        vector_db = MilvusVDB(
+            collection_name=settings.collection_name,
+            vector_dim=settings.vector_dim,
+            batch_size=settings.batch_size,
+            db_file=settings.milvus_db_path
+        )
+        
+        # Initialize database and create collection
+        vector_db.initialize_client()
+        vector_db.create_collection()
+        
+        # Ingest data
+        vector_db.ingest_data(embed_data)
+        logger.info("Vector database setup completed")
+        
+        # Step 4: Setup retrieval system
+        logger.info("Step 4: Setting up retrieval system...")
+        
+        retriever = Retriever(
+            vector_db=vector_db,
+            embed_data=embed_data,
+            top_k=settings.top_k
+        )
+        
+        # Step 5: Setup RAG system
+        logger.info("Step 5: Setting up RAG system...")
+        
+        rag_system = RAG(
+            retriever=retriever,
+            llm_model=settings.llm_model,
+            openai_api_key=settings.openai_api_key,
+            temperature=settings.temperature,
+            max_tokens=settings.max_tokens
+        )
+        
+        # Step 6: Setup workflow
+        logger.info("Step 6: Setting up agentic workflow...")
+        
+        workflow = ParalegalAgentWorkflow(
+            retriever=retriever,
+            rag_system=rag_system,
+            firecrawl_api_key=os.getenv("FIRECRAWL_API_KEY"),
+            openai_api_key=os.getenv("OPENAI_API_KEY")
+        )
+        
+        logger.info("Setup completed! Ready for queries.")
+        
+        print("\n" + "="*60)
+        print("Pipeline Ready!")
+        print("Type 'quit' to exit, 'help' for commands")
+        print("="*60)
+        
+        while True:
+            try:
+                query = input("\nEnter your question: ").strip()
+                
+                if query.lower() in ['quit', 'exit', 'q']:
+                    break
+                elif not query:
+                    continue
+                
+                logger.info(f"Processing query: {query}")
+                
+                # Run the workflow
+                result = await workflow.run_workflow(query)
+                
+                # Display results
+                print("\n" + "-"*60)
+                print("ANSWER:")
+                print(result["answer"])
+                
+                if result.get("web_search_used", False):
+                    print(f"\n🌐 Web search was used to enhance the response")
+                else:
+                    print(f"\n📚 Response based on document knowledge")
+                    # Show citations grounding the answer
+                    try:
+                        citations = get_citations(retriever, query, top_k=min(3, settings.top_k))
+                        print_citations(citations)
+                    except Exception as e:
+                        logger.warning(f"Could not fetch citations: {e}")
+                
+                print("-"*60)
+                
+                show_details = input("\nShow detailed results? (y/n): ").strip().lower()
+                if show_details == 'y':
+                    print_detailed_results(result)
+                
+            except KeyboardInterrupt:
+                print("\nExiting...")
+                break
+            except Exception as e:
+                logger.error(f"Error processing query: {e}")
+                print(f"Error: {e}")
+        
+        # Cleanup
+        logger.info("Cleaning up...")
+        vector_db.close()
+        logger.info("Demo completed")
+        
+    except Exception as e:
+        logger.error(f"Pipeline setup failed: {e}")
+        print(f"Setup failed: {e}")
+
+def print_detailed_results(result):
+    print("\n" + "="*60)
+    print("DETAILED RESULTS")
+    print("="*60)
+    
+    print(f"\nOriginal Query: {result['query']}")
+    
+    if result.get('rag_response'):
+        print(f"\nRAG Response:")
+        print(result['rag_response'])
+    
+    if result.get('web_search_used') and result.get('web_results'):
+        print(f"\nWeb Search Results:")
+        print(result['web_results'][:500] + "..." if len(result['web_results']) > 500 else result['web_results'])
+    
+    if result.get('error'):
+        print(f"\nError: {result['error']}")
+    
+    print("="*60)
+
+async def test_retrieval():
+    # Test retrieval pipeline
+    logger.info("Running retrieval test...")
+
+    test_text = [
+        "This is a test document about artificial intelligence.",
+        "Machine learning is a subset of artificial intelligence.",
+        "Deep learning uses neural networks with multiple layers."
+    ]
+    
+    # Test embedding
+    embed_data = EmbedData()
+    embed_data.embed(test_text)
+    
+    # Test vector database
+    vector_db = MilvusVDB(collection_name="test_collection")
+    vector_db.initialize_client()
+    vector_db.create_collection()
+    vector_db.ingest_data(embed_data)
+    
+    # Test retrieval
+    retriever = Retriever(vector_db, embed_data)
+    results = retriever.search("What is machine learning?")
+    
+    logger.info(f"Test completed. Retrieved {len(results)} results")
+
+    # Test citations
+    citations = get_citations(retriever, "What is machine learning?", top_k=3)
+    print(citations)
+    
+    # Cleanup
+    vector_db.close()
+    
+    return True
+
+if __name__ == "__main__":
+    asyncio.run(main())

paralegal-agent-crew/pyproject.toml

@@ -0,0 +1,23 @@
+[project]
+name = "paralegal-agent"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "firecrawl-py>=0.0.16",
+    "crewai>=0.74.0",
+    "crewai-tools>=0.10.0",
+    "loguru>=0.7.0",
+    "nest-asyncio>=1.6.0",
+    "numpy>=1.24.0",
+    "pydantic>=2.0.0",
+    "pydantic-settings>=2.10.1",
+    "pymilvus>=2.4.0",
+    "python-dotenv>=1.0.0",
+    "streamlit>=1.48.0",
+    "sentence-transformers>=3.0.0",
+    "pypdf>=4.2.0",
+    "openai>=1.43.0",
+    "ansi2html>=1.9.1",
+]

paralegal-agent-crew/src/embeddings/embed_data.py

@@ -0,0 +1,93 @@
+import numpy as np
+from typing import List
+from loguru import logger
+from sentence_transformers import SentenceTransformer
+from config.settings import settings
+
+def batch_iterate(lst: List, batch_size: int):
+    for i in range(0, len(lst), batch_size):
+        yield lst[i:i+batch_size]
+
+class EmbedData:
+    """Handles document embedding with binary quantization support."""
+    def __init__(
+        self, 
+        embed_model_name: str = None,
+        batch_size: int = None,
+        cache_folder: str = None
+    ):
+        self.embed_model_name = embed_model_name or settings.embedding_model
+        self.batch_size = batch_size or settings.batch_size
+        self.cache_folder = cache_folder or settings.hf_cache_dir
+        
+        self.embed_model = self._load_embed_model()
+        self.embeddings = []
+        self.binary_embeddings = []
+        self.contexts = []
+
+    def _load_embed_model(self):
+        """Load the embedding model using sentence-transformers"""
+        logger.info(f"Loading embedding model: {self.embed_model_name}")
+        model = SentenceTransformer(
+            model_name_or_path=self.embed_model_name,
+            cache_folder=self.cache_folder,
+            trust_remote_code=True,
+        )
+        return model
+
+    def _binary_quantize(self, embeddings: List[List[float]]):
+        """Convert float32 embeddings to binary vectors."""
+        embeddings_array = np.array(embeddings)
+        binary_embeddings = np.where(embeddings_array > 0, 1, 0).astype(np.uint8)
+        
+        # Pack bits into bytes (8 dimensions per byte)
+        packed_embeddings = np.packbits(binary_embeddings, axis=1)
+        return [vec.tobytes() for vec in packed_embeddings]
+
+    def generate_embedding(self, contexts: List[str]):
+        embeddings = self.embed_model.encode(
+            sentences=contexts,
+            batch_size=min(self.batch_size, max(1, len(contexts))),
+            convert_to_numpy=True,
+            normalize_embeddings=False,
+            show_progress_bar=False,
+        )
+        return embeddings.tolist()
+
+    def embed(self, contexts: List[str]):
+        self.contexts = contexts
+        logger.info(f"Generating embeddings for {len(contexts)} contexts...")
+
+        for batch_context in batch_iterate(contexts, self.batch_size):
+            # Generate float32 embeddings
+            batch_embeddings = self.generate_embedding(batch_context)
+            self.embeddings.extend(batch_embeddings)
+
+            # Convert to binary and store
+            binary_batch = self._binary_quantize(batch_embeddings)
+            self.binary_embeddings.extend(binary_batch)
+
+        logger.info(f"Generated {len(self.embeddings)} embeddings with binary quantization")
+
+    def get_query_embedding(self, query: str):
+        # Generate embedding for a single query
+        embedding = self.embed_model.encode(
+            sentences=[query],
+            convert_to_numpy=True,
+            normalize_embeddings=False,
+            show_progress_bar=False,
+        )
+        return embedding[0].tolist()
+
+    def binary_quantize_query(self, query_embedding: List[float]):
+        # Convert query embedding to binary format
+        embedding_array = np.array([query_embedding])
+        binary_embedding = np.where(embedding_array > 0, 1, 0).astype(np.uint8)
+        packed_embedding = np.packbits(binary_embedding, axis=1)
+        return packed_embedding[0].tobytes()
+
+    def clear(self):
+        self.embeddings.clear()
+        self.binary_embeddings.clear()
+        self.contexts.clear()
+        logger.info("Cleared all embeddings and contexts")

paralegal-agent-crew/src/generation/rag.py

@@ -0,0 +1,98 @@
+from typing import Optional
+from loguru import logger
+from crewai import LLM
+from pydantic import BaseModel
+from src.retrieval.retriever_rerank import Retriever
+from config.settings import settings
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class RAG: 
+    def __init__(
+        self, 
+        retriever: Retriever, 
+        llm_model: str = None,
+        openai_api_key: str = None,
+        temperature: float = None,
+        max_tokens: int = None
+    ):
+        self.retriever = retriever
+        self.llm_model = llm_model or settings.llm_model
+        self.openai_api_key = openai_api_key or settings.openai_api_key
+        self.temperature = temperature or settings.temperature
+        self.max_tokens = max_tokens or settings.max_tokens
+        
+        # Initialize LLM
+        self.llm = self._setup_llm()
+        
+        # System message
+        self.system_message = ChatMessage(
+            role="system",
+            content="You are a helpful assistant that answers questions based on the provided context. "
+                   "Always base your answers on the given information and clearly indicate when you don't know something."
+        )
+        
+        # RAG prompt template
+        self.prompt_template = (
+            "CONTEXT:\n"
+            "{context}\n"
+            "---------------------\n"
+            "Based on the context information above, please answer the following question. "
+            "If the context doesn't contain enough information to answer the question, or "
+            "even if you know the answer, but it is not relevant to the provided context, "
+            "clearly state that you don't know and explain what information is missing.\n\n"
+            "QUESTION: {query}\n"
+            "ANSWER: "
+        )
+
+    def _setup_llm(self):
+        if not self.openai_api_key:
+            raise ValueError(
+                "OpenAI API key is required. Set OPENAI_API_KEY environment variable "
+                "or pass openai_api_key parameter."
+            )
+        llm = LLM(model=self.llm_model, api_key=self.openai_api_key, temperature=self.temperature)
+        logger.info(f"Initialized CrewAI LLM with model: {self.llm_model}")
+        return llm
+
+    def generate_context(self, query: str, top_k: Optional[int] = None):
+        # Generate context from retrieval results
+        return self.retriever.get_combined_context(query, top_k)
+
+    def query(self, query: str, top_k: Optional[int] = None):
+        # Generate context from retrieval
+        context = self.generate_context(query, top_k)
+        
+        # Create prompt from template
+        prompt = self.prompt_template.format(context=context, query=query)
+        return self.llm.call(f"{self.system_message.content}\n\n{prompt}")
+
+    def get_detailed_response(self, query: str, top_k: Optional[int] = None):
+        # Get retrieval results with scores
+        retrieval_results = self.retriever.search_with_scores(query, top_k)
+        
+        # Generate context
+        context = self.retriever.get_combined_context(query, top_k)
+        
+        # Generate response
+        response = self.query(query, top_k=top_k)
+        
+        return {
+            "response": response,
+            "context": context,
+            "sources": retrieval_results,
+            "query": query,
+            "model": self.llm_model
+        }
+
+    def set_prompt_template(self, template: str):
+        # Set custom prompt template
+        self.prompt_template = template
+        logger.info("Updated prompt template")
+
+    def set_system_message(self, content: str):
+        # Set custom system message
+        self.system_message = ChatMessage(role="system", content=content)
+        logger.info("Updated system message")

paralegal-agent-crew/src/indexing/milvus_vdb.py

@@ -0,0 +1,170 @@
+from typing import List
+from loguru import logger
+from pymilvus import MilvusClient, DataType
+from src.embeddings.embed_data import EmbedData, batch_iterate
+from config.settings import settings
+
+class MilvusVDB:
+    """Milvus vector database with binary quantization support."""
+    def __init__(
+        self,
+        collection_name: str = None,
+        vector_dim: int = None,
+        batch_size: int = None,
+        db_file: str = None
+    ):
+        self.collection_name = collection_name or settings.collection_name
+        self.vector_dim = vector_dim or settings.vector_dim
+        self.batch_size = batch_size or settings.batch_size
+        self.db_file = db_file or settings.milvus_db_path
+        self.client = None
+
+    def initialize_client(self):
+        try:
+            self.client = MilvusClient(self.db_file)
+            logger.info(f"Initialized Milvus Lite client with database: {self.db_file}")
+        except Exception as e:
+            logger.error(f"Failed to initialize Milvus client: {e}")
+            raise e
+
+    def create_collection(self):
+        """Create collection with binary vector support."""
+        if not self.client:
+            raise RuntimeError("Milvus client not initialized. Call initialize_client() first.")
+        
+        # Drop existing collection if it exists
+        if self.client.has_collection(collection_name=self.collection_name):
+            self.client.drop_collection(collection_name=self.collection_name)
+            logger.info(f"Dropped existing collection: {self.collection_name}")
+
+        # Create schema for binary vectors
+        schema = self.client.create_schema(
+            auto_id=True,
+            enable_dynamic_fields=True,
+        )
+
+        # Add fields to schema
+        schema.add_field(
+            field_name="id", 
+            datatype=DataType.INT64, 
+            is_primary=True, 
+            auto_id=True
+        )
+        schema.add_field(
+            field_name="context", 
+            datatype=DataType.VARCHAR, 
+            max_length=65535
+        )
+        schema.add_field(
+            field_name="binary_vector", 
+            datatype=DataType.BINARY_VECTOR, 
+            dim=self.vector_dim
+        )
+
+        # Create index parameters for binary vectors
+        index_params = self.client.prepare_index_params()
+        index_params.add_index(
+            field_name="binary_vector",
+            index_name="binary_vector_index",
+            index_type="BIN_FLAT",  # Exact search for binary vectors
+            metric_type="HAMMING"   # Hamming distance for binary vectors
+        )
+
+        # Create collection with schema and index
+        self.client.create_collection(
+            collection_name=self.collection_name,
+            schema=schema,
+            index_params=index_params
+        )
+
+        logger.info(f"Created collection '{self.collection_name}' with binary vectors (dim={self.vector_dim})")
+
+    def ingest_data(self, embed_data: EmbedData):
+        """Ingest embedded data into the vector database."""
+        if not self.client:
+            raise RuntimeError("Milvus client not initialized. Call initialize_client() first.")
+        
+        logger.info(f"Ingesting {len(embed_data.contexts)} documents...")
+
+        total_inserted = 0
+        for batch_context, batch_binary_embeddings in zip(
+            batch_iterate(embed_data.contexts, self.batch_size),
+            batch_iterate(embed_data.binary_embeddings, self.batch_size)
+        ):
+            # Prepare data for insertion
+            data_batch = []
+            for context, binary_embedding in zip(batch_context, batch_binary_embeddings):
+                data_batch.append({
+                    "context": context,
+                    "binary_vector": binary_embedding
+                })
+
+            # Insert batch
+            self.client.insert(
+                collection_name=self.collection_name,
+                data=data_batch
+            )
+
+            total_inserted += len(batch_context)
+            logger.info(f"Inserted batch: {len(batch_context)} documents")
+
+        logger.info(f"Successfully ingested {total_inserted} documents with binary quantization")
+
+    def search(
+        self, 
+        binary_query: bytes, 
+        top_k: int = None,
+        output_fields: List[str] = None
+    ):
+        if not self.client:
+            raise RuntimeError("Milvus client not initialized. Call initialize_client() first.")
+        
+        top_k = top_k or settings.top_k
+        output_fields = output_fields or ["context"]
+
+        # Perform similarity search using MilvusClient
+        search_results = self.client.search(
+            collection_name=self.collection_name,
+            data=[binary_query],
+            anns_field="binary_vector",
+            search_params={"metric_type": "HAMMING", "params": {}},
+            limit=top_k,
+            output_fields=output_fields
+        )
+
+        # Format results
+        formatted_results = []
+        for result in search_results[0]:
+            formatted_results.append({
+                "id": result["id"],
+                "score": 1.0 / (1.0 + result["distance"]),  # Convert Hamming distance to similarity
+                "payload": {"context": result["entity"]["context"]}
+            })
+
+        return formatted_results
+
+    def collection_exists(self):
+        if not self.client:
+            return False
+        return self.client.has_collection(collection_name=self.collection_name)
+
+    def get_collection_info(self):
+        if not self.client:
+            raise RuntimeError("Milvus client not initialized. Call initialize_client() first.")
+        
+        if not self.collection_exists():
+            return {"exists": False}
+        
+        # Get collection statistics
+        stats = self.client.get_collection_stats(collection_name=self.collection_name)
+        return {
+            "exists": True,
+            "row_count": stats["row_count"],
+            "collection_name": self.collection_name
+        }
+
+    # Close the database connection
+    def close(self):
+        if self.client:
+            self.client.close()
+            logger.info("Closed Milvus client connection")

paralegal-agent-crew/src/retrieval/retriever_rerank.py

@@ -0,0 +1,102 @@
+from typing import Optional
+from loguru import logger
+from pydantic import BaseModel, Field
+from src.indexing.milvus_vdb import MilvusVDB
+from src.embeddings.embed_data import EmbedData
+from config.settings import settings
+
+class TextNode(BaseModel):
+    text: str
+    id_: str
+    metadata: dict | None = Field(default=None)
+
+class NodeWithScore(BaseModel):
+    node: TextNode
+    score: float
+
+class Retriever:
+    def __init__(
+        self, 
+        vector_db: MilvusVDB, 
+        embed_data: EmbedData, 
+        top_k: int = None
+    ):
+        self.vector_db = vector_db
+        self.embed_data = embed_data
+        self.top_k = top_k or settings.top_k
+
+    def search(self, query: str, top_k: Optional[int] = None):
+        """Search for relevant documents using vector similarity."""
+        if top_k is None:
+            top_k = self.top_k
+
+        # Generate query embedding and convert to binary
+        query_embedding = self.embed_data.get_query_embedding(query)
+        binary_query = self.embed_data.binary_quantize_query(query_embedding)
+
+        # Perform vector search
+        search_results = self.vector_db.search(
+            binary_query=binary_query,
+            top_k=top_k,
+            output_fields=["context"]
+        )
+
+        # Convert to NodeWithScore objects
+        nodes_with_scores = []
+        for result in search_results:
+            node = TextNode(
+                text=result["payload"]["context"],
+                id_=str(result["id"])
+            )
+            node_with_score = NodeWithScore(
+                node=node,
+                score=result["score"]
+            )
+            nodes_with_scores.append(node_with_score)
+
+        # logger.info(f"Retrieved {len(nodes_with_scores)} results for query")
+        return nodes_with_scores
+
+    def get_contexts(self, query: str, top_k: Optional[int] = None):
+        nodes_with_scores = self.search(query, top_k)
+        return [node.node.text for node in nodes_with_scores]
+
+    def get_combined_context(self, query: str, top_k: Optional[int] = None):
+        contexts = self.get_contexts(query, top_k)
+        return "\n\n---\n\n".join(contexts)
+
+    def search_with_scores(self, query: str, top_k: Optional[int] = None):
+        nodes_with_scores = self.search(query, top_k)
+        
+        results = []
+        for node_with_score in nodes_with_scores:
+            results.append({
+                "context": node_with_score.node.text,
+                "score": node_with_score.score,
+                "node_id": node_with_score.node.id_,
+                "metadata": node_with_score.node.metadata or {}
+            })
+        
+        return results
+
+    def get_citations(self, query: str, top_k: int = 3, snippet_chars: int = 300):
+        # Return top-k retrieval results formatted as citation dicts
+        results = self.search_with_scores(query, top_k)
+        citations = []
+        for rank, item in enumerate(results, start=1):
+            context = (item.get("context") or "").strip()
+            if context:
+                snippet = context[:snippet_chars]
+                if len(context) > snippet_chars:
+                    snippet += "…"
+            else:
+                snippet = ""
+
+            citations.append({
+                "rank": rank,
+                "node_id": item.get("node_id"),
+                "score": item.get("score"),
+                "snippet": snippet,
+                "metadata": item.get("metadata") or {},
+            })
+        return citations

paralegal-agent-crew/src/tools/firecrawl_search_tool.py

@@ -0,0 +1,48 @@
+from typing import Type
+from pydantic import BaseModel, Field
+from crewai.tools import BaseTool
+from firecrawl import FirecrawlApp
+from config.settings import settings
+import os
+
+class FirecrawlSearchInput(BaseModel):
+    """Input schema for Firecrawl web search tool"""
+    query: str = Field(..., description="The search query to look up on the web.")
+    limit: int = Field(..., description="Maximum number of results to fetch.")
+
+class FirecrawlSearchTool(BaseTool):
+    name: str = "Firecrawl Web Search"
+    description: str = (
+        "Search the web using Firecrawl and return a concise list of results "
+        "(title, URL, and short description snippet)."
+    )
+    args_schema: Type[BaseModel] = FirecrawlSearchInput
+
+    def _run(self, query: str, limit: int = 3) -> str:
+        api_key = settings.firecrawl_api_key or os.getenv("FIRECRAWL_API_KEY")
+        if not api_key:
+            return "Web search unavailable - API not configured."
+
+        try:
+            app = FirecrawlApp(api_key=api_key)
+            response = app.search(query, limit=limit)
+            results_list = getattr(response, "data", None)
+
+            if not isinstance(results_list, list) or not results_list:
+                return "No relevant web search results found."
+
+            search_contents= []
+            for result in results_list:
+                if not isinstance(result, dict):
+                    continue
+                url = result.get("url", "No URL")
+                title = result.get("title", "No title")
+                description = (result.get("description") or "").strip()
+                snippet = description[:1000] if description else "[no description available]"
+                search_contents.append(f"Title: {title}\nURL: {url}\nContent: {snippet}")
+
+            return "\n\n---\n\n".join(search_contents) if search_contents else "No relevant web search results found."
+        except Exception as e:
+            return f"Web search unavailable due to technical issues: {e}"
+
+

paralegal-agent-crew/src/workflows/agent_workflow.py

@@ -0,0 +1,234 @@
+from typing import Optional, Any
+from loguru import logger
+from crewai import LLM
+from crewai.flow.flow import Flow, start, listen, router, or_
+from pydantic import BaseModel
+
+from .events import RetrieveEvent, EvaluateEvent, WebSearchEvent, SynthesizeEvent
+from src.tools.firecrawl_search_tool import FirecrawlSearchTool
+from src.retrieval.retriever_rerank import Retriever
+from src.generation.rag import RAG
+from config.settings import settings
+
+# Prompt templates for workflow steps
+ROUTER_EVALUATION_TEMPLATE = (
+    """You are a quality evaluator for RAG responses. Your task is to determine if the given response adequately answers the user's question.
+
+USER QUESTION:
+{query}
+
+RAG RESPONSE:
+{rag_response}
+
+EVALUATION CRITERIA:
+- Does the response directly address the user's question?
+- Is the response factually coherent and well-structured?
+- Does the response contain sufficient detail to be helpful?
+- If the response says "I don't know" or similar, is it because the context truly lacks the information?
+
+Please evaluate the response quality and respond with either:
+- "GOOD" - if the response adequately answers the question
+- "BAD" - if the response is incomplete, unclear, or doesn't answer the question
+
+IMPORTANT: Respond with ONLY ONE WORD in UPPERCASE: GOOD or BAD. No punctuation or extra text.
+
+Your evaluation (GOOD or BAD):"""
+)
+
+QUERY_OPTIMIZATION_TEMPLATE = (
+    """Optimize the following query for web search to get the most relevant and accurate results.
+
+Original Query: {query}
+
+Guidelines:
+- Make the query more specific and searchable
+- Add relevant keywords that would help find authoritative sources
+- Keep it concise but comprehensive
+- Focus on the core information need
+
+Optimized Query:"""
+)
+
+SYNTHESIS_TEMPLATE = (
+    """You are a response synthesizer. Create a comprehensive and accurate answer based on the available information.
+
+USER QUESTION:
+{query}
+
+RAG RESPONSE (from document knowledge):
+{rag_response}
+
+WEB SEARCH RESULTS (additional context):
+{web_results}
+
+INSTRUCTIONS:
+- Synthesize information from both sources to provide the most complete answer
+- Prioritize information from reliable sources
+- If there are contradictions, acknowledge them
+- Clearly indicate when information comes from web search vs document knowledge
+- If web results are empty, refine and improve the RAG response
+
+SYNTHESIZED RESPONSE:"""
+)
+
+# Define flow state
+class ParalegalAgentState(BaseModel):
+    query: str = ""
+    top_k: Optional[int] = 3
+
+class ParalegalAgentWorkflow(Flow[ParalegalAgentState]):
+    """Paralegal Agent Workflow with router and web search fallback using CrewAI Flows."""
+
+    def __init__(
+        self,
+        retriever: Retriever,
+        rag_system: RAG,
+        firecrawl_api_key: Optional[str] = None,
+        openai_api_key: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.retriever = retriever
+        self.rag = rag_system
+
+        self.openai_api_key = openai_api_key or settings.openai_api_key
+        self.llm = LLM(model=settings.llm_model, api_key=self.openai_api_key, temperature=0.1)
+
+    @start()
+    def retrieve(self) -> RetrieveEvent:
+        """Retrieve relevant documents from vector database"""
+        query = self.state.query
+        top_k = self.state.top_k
+
+        if not query:
+            raise ValueError("Query is required")
+
+        logger.info(f"Retrieving documents for query: {query}")
+
+        retrieved_nodes = self.retriever.search(query, top_k=top_k)
+        logger.info(f"Retrieved {len(retrieved_nodes)} documents")
+        return RetrieveEvent(retrieved_nodes=retrieved_nodes, query=query)
+
+    @listen(retrieve)
+    def generate_rag_response(self, ev: RetrieveEvent) -> EvaluateEvent:
+        """Generate initial RAG response"""
+        query = ev.query
+        retrieved_nodes = ev.retrieved_nodes
+
+        logger.info("Generating RAG response")
+
+        rag_response = self.rag.query(query)
+
+        logger.info("RAG response generated")
+        return EvaluateEvent(
+            rag_response=rag_response,
+            retrieved_nodes=retrieved_nodes,
+            query=query
+        )
+
+    @router(generate_rag_response)
+    def evaluate_response(self, ev: EvaluateEvent) -> str:
+        """Evaluate RAG response quality and route accordingly"""
+        rag_response = ev.rag_response
+        query = ev.query
+        
+        logger.info("Evaluating RAG response quality")
+
+        evaluation_prompt = ROUTER_EVALUATION_TEMPLATE.format(query=query, rag_response=rag_response)
+        resp_text = self.llm.call(evaluation_prompt)
+        evaluation = (resp_text or "").strip().upper().split()[0]
+
+        logger.info(f"Evaluation result: {evaluation}")
+        return "synthesize" if "GOOD" in evaluation else "web_search"
+
+    @listen("web_search")
+    def perform_web_search(self, ev: EvaluateEvent | WebSearchEvent) -> SynthesizeEvent:
+        """Perform web search if insufficient information from RAG response"""
+        query = ev.query
+        rag_response = ev.rag_response
+        retrieved_nodes = getattr(ev, "retrieved_nodes", [])
+        
+        logger.info("Performing web search")
+        
+        search_results = ""
+        try:
+            optimization_prompt = QUERY_OPTIMIZATION_TEMPLATE.format(query=query)
+            optimized_query = (self.llm.call(optimization_prompt) or query).strip()
+            search_results = FirecrawlSearchTool().run(query=optimized_query, limit=3)
+            logger.info("Web search completed via custom tool")
+        except Exception as e:
+            logger.error(f"Web search failed: {e}")
+            search_results = "Web search unavailable due to technical issues."
+        
+        return SynthesizeEvent(
+            rag_response=rag_response,
+            web_search_results=search_results,
+            retrieved_nodes=retrieved_nodes,
+            query=query,
+            use_web_results=True
+        )
+
+    @listen(or_("synthesize", "perform_web_search"))
+    def synthesize_response(self, ev: EvaluateEvent | SynthesizeEvent) -> dict:
+        """Synthesize final response from RAG and web search results"""
+        rag_response = ev.rag_response
+        web_results = getattr(ev, "web_search_results", "") or ""
+        query = ev.query
+        use_web_results = getattr(ev, "use_web_results", False)
+        
+        logger.info("Synthesizing final response")
+        
+        if use_web_results and web_results:
+            synthesis_prompt = SYNTHESIS_TEMPLATE.format(
+                query=query, rag_response=rag_response, web_results=web_results
+            )
+            synthesized_answer = self.llm.call(synthesis_prompt)
+            result = {
+                "answer": synthesized_answer,
+                "rag_response": rag_response,
+                "web_search_used": True,
+                "web_results": web_results,
+                "query": query,
+            }
+        else:
+            refinement_prompt = (
+                f"Improve and refine the following response to make it more helpful and comprehensive:\n\n"
+                f"Original Response: {rag_response}\n\nRefined Response:"
+            )
+            refined = self.llm.call(refinement_prompt)
+            result = {
+                "answer": refined,
+                "rag_response": rag_response,
+                "web_search_used": False,
+                "web_results": None,
+                "query": query,
+            }
+        
+        logger.info("Final response synthesized")
+        return result
+
+    async def run_workflow(self, query: str, top_k: Optional[int] = None) -> dict:
+        """
+        Run the complete flow for a given query.
+        
+        Args:
+            query: User question
+            top_k: Number of documents to retrieve
+            
+        Returns:
+            Dictionary with final answer and metadata
+        """
+        try:
+            # Kick off the CrewAI flow asynchronously with runtime inputs
+            result = await self.kickoff_async(inputs={"query": query, "top_k": top_k})
+            return result if isinstance(result, dict) else {"answer": str(result), "query": query}
+        except Exception as e:
+            logger.error(f"Workflow execution failed: {e}")
+            return {
+                "answer": f"I apologize, but I encountered an error while processing your question: {str(e)}",
+                "rag_response": None,
+                "web_search_used": False,
+                "web_results": None,
+                "query": query,
+                "error": str(e)
+            }

paralegal-agent-crew/src/workflows/events.py

@@ -0,0 +1,28 @@
+from typing import List, Optional
+from pydantic import BaseModel
+from src.retrieval.retriever_rerank import NodeWithScore
+
+class RetrieveEvent(BaseModel):
+    """Event containing retrieved nodes from vector database."""
+    retrieved_nodes: List[NodeWithScore]
+    query: str
+
+class EvaluateEvent(BaseModel): 
+    """Event for evaluating RAG response quality."""
+    rag_response: str
+    retrieved_nodes: List[NodeWithScore]
+    query: str
+
+class WebSearchEvent(BaseModel):
+    """Event for web search when RAG response is insufficient."""
+    rag_response: str
+    query: str
+    search_results: Optional[str] = None
+
+class SynthesizeEvent(BaseModel):
+    """Event for final response synthesis."""
+    rag_response: str
+    retrieved_nodes: List[NodeWithScore]
+    query: str
+    web_search_results: Optional[str] = None
+    use_web_results: bool = False