DEV Community

James Li
James Li

Posted on

1 1 1 1 1

Multi-Agent Hybrid Knowledge Base Retrieval: Building a High-Precision Legal Case Analysis Platform

This article explores how to leverage LangGraph's multi-agent architecture and hybrid knowledge base technology to build a high-precision legal case analysis platform, focusing on the Enron litigation case document analysis scenario, demonstrating how the combination of GraphRAG, Neo4j, and vector databases can solve complex challenges in legal document retrieval.

1. Technical Challenges in Legal RAG Systems

1.1 Business Requirements Analysis

According to our client requirements, they need a legal case file search application that primarily processes Enron case files released by the U.S. Department of Justice. These files contain numerous emails, financial statements, legal documents, and other multi-format documents.

Core business requirements include:

  • Processing multi-format, multi-language legal documents
  • Implementing complex multi-dimensional queries (time, relationships between individuals, events)
  • Providing high-precision retrieval results and complete evidence chains
  • Supporting entity relationship visualization and timeline analysis

1.2 Technical Challenge Analysis

Compared to ordinary RAG systems, legal RAG faces the following unique challenges:

  • Complex Entity Relationships: Need to accurately identify and model complex associations between people, organizations, and events
  • Importance of Temporal Relationships: The sequence of events in legal cases is critical and affects causal judgments
  • Multi-hop Reasoning Requirements: Need to build complete evidence chains across multiple documents
  • High Query Complexity: Queries like "When did A first mention financial problems at Company C to B?" require multi-dimensional analysis

2. System Architecture Design

2.1 Overall Architecture

Based on these requirements, we designed a multi-agent hybrid knowledge base system based on the LangGraph Supervisor architecture:

Image description

2.2 Hybrid Knowledge Base Design

We employed a combination of three different types of knowledge bases:

  1. Neo4j Graph Database: Stores entity relationship networks

    • Node types: People, organizations, documents, events
    • Relationship types: Sent email, attended meeting, belongs to, mentioned, etc.
  2. Weaviate Vector Database: Stores document semantic representations

    • Document-level vectors: Semantic representation of entire documents
    • Paragraph-level vectors: Semantic representation of fine-grained text blocks
  3. Relational Database: Stores structured metadata

    • Document metadata: Time, author, type, etc.
    • Entity attributes: Position, role, time range, etc.

3. Core Implementation Code

3.1 LangGraph Supervisor Architecture Core

from langgraph.graph import StateGraph, MessagesState, START, END
from typing import Literal, List, Dict, Any

# Define system state
class LegalCaseState(MessagesState):
    """Legal case analysis system state"""
    next: str
    entities: List[Dict[str, Any]] = []
    retrieved_documents: List[Dict[str, Any]] = []
    evidence_chain: List[Dict[str, Any]] = []

# Supervisor logic - core decision point of the system
def supervisor(state: LegalCaseState):
    system_prompt = (
        "You are a legal case analysis supervisor managing multiple agents.\n"
        "Given the current state and user query, decide which agent should act next."
    )

    messages = [{"role": "system", "content": system_prompt}] + state["messages"]
    response = llm.with_structured_output(Router).invoke(messages)

    next_ = response["next"]
    if next_ == "FINISH":
        next_ = END

    return {"next": next_}

# Build graph
builder = StateGraph(LegalCaseState)
builder.add_node("supervisor", supervisor)

# Add conditional edges - dynamic decision flow
builder.add_conditional_edges("supervisor", lambda state: state["next"])
builder.add_edge(START, "supervisor")
Enter fullscreen mode Exit fullscreen mode

3.2 Hybrid Retrieval Strategy Implementation

def hybrid_retrieval(query, top_k=5):
    """Execute hybrid retrieval strategy - core innovation point of the system"""

    # 1. Query parsing - extract key entities and relationships
    parsed_query = parse_query(query)

    # 2. Vector retrieval - semantic similarity
    query_vector = embed_text(query)
    vector_results = weaviate_client.query.get(
        "LegalDocument", ["content", "title", "date", "author"]
    ).with_near_vector({
        "vector": query_vector
    }).with_limit(top_k).do()

    # 3. Graph retrieval - structured relationships
    cypher_query = build_cypher_query(parsed_query)
    graph_results = neo4j_client.query(cypher_query) if cypher_query else []

    # 4. Retrieval enhancement - extract relevant document IDs from graph retrieval results
    doc_ids_from_graph = extract_document_ids(graph_results)

    # 5. Result merging and reranking - key innovation point
    combined_results = merge_and_rank_results(
        vector_results["data"]["Get"]["LegalDocument"],
        get_documents_by_ids(doc_ids_from_graph),
        graph_results
    )

    return combined_results
Enter fullscreen mode Exit fullscreen mode

3.3 Reasoning Agent Implementation

def reasoner(state: LegalCaseState):
    """Reason based on retrieval results, build evidence chain - core analytical capability"""
    query = state["messages"][0].content
    graph_results = next((m.content for m in state["messages"] 
                         if getattr(m, "name", "") == "graph_retriever"), "")
    vector_results = next((m.content for m in state["messages"] 
                          if getattr(m, "name", "") == "vector_retriever"), "")

    # Combine retrieval results
    combined_context = f"Graph Results:\n{graph_results}\n\nVector Results:\n{vector_results}"

    # Build reasoning prompt
    reasoning_prompt = f"""
    Based on the retrieved information, please:
    1. Identify key evidence relevant to the query
    2. Establish connections between evidence pieces
    3. Build a logical evidence chain
    4. Identify any contradictions
    5. Provide a well-reasoned answer

    Query: {query}
    Retrieved Information: {combined_context}
    """

    # Execute reasoning
    reasoning_result = llm.invoke(reasoning_prompt)

    # Extract evidence chain - structured output
    evidence_chain = extract_evidence_chain(reasoning_result.content)

    return {
        "messages": [HumanMessage(content=reasoning_result.content)],
        "evidence_chain": evidence_chain
    }
Enter fullscreen mode Exit fullscreen mode

3.4 Entity Disambiguation Processing - Key Technical Challenge

def entity_disambiguation(entities):
    """Entity disambiguation processing - solving entity recognition challenges in legal documents"""

    # 1. Build entity index
    entity_index = {}
    for entity in entities:
        key = normalize_name(entity["name"])
        if key not in entity_index:
            entity_index[key] = []
        entity_index[key].append(entity)

    # 2. Merge similar entities - core algorithm
    merged_entities = []
    for key, entity_group in entity_index.items():
        if len(entity_group) > 1:
            # Use context and attributes for disambiguation
            similarity_matrix = calculate_entity_similarity(entity_group)
            clusters = cluster_entities(entity_group, similarity_matrix, threshold=0.85)

            # Merge entities in each cluster
            for cluster in clusters:
                merged_entity = merge_entity_cluster(cluster)
                merged_entities.append(merged_entity)
        else:
            merged_entities.extend(entity_group)

    return merged_entities
Enter fullscreen mode Exit fullscreen mode

3.5 Efficient Neo4j Batch Import

def import_entities_to_neo4j(entities, relations):
    """Efficient batch import to Neo4j - key performance optimization point"""

    # Batch import entities - use parameterized queries and batching to improve performance
    batch_size = 1000
    for i in range(0, len(entities), batch_size):
        batch = entities[i:i+batch_size]
        query = """
        UNWIND $batch AS entity
        MERGE (n:`${label}` {id: entity.id})
        SET n += entity.properties
        """

        neo4j_client.query(
            query,
            parameters={"batch": [
                {
                    "id": e["id"],
                    "label": e["type"],
                    "properties": e["properties"]
                } for e in batch
            ]}
        )
Enter fullscreen mode Exit fullscreen mode

4. System Validation

4.1 Query Examples and Results

Here are examples of how the system processes complex queries:

Query 1: "What communications occurred between Andrew Fastow and Kenneth Lay regarding Special Purpose Entities (SPE) from July to August 2001?"

result = graph.invoke({
    "messages": "What communications occurred between Andrew Fastow and Kenneth Lay regarding Special Purpose Entities (SPE) from July to August 2001?"
})

print(result["messages"][-1].content)
Enter fullscreen mode Exit fullscreen mode

Output:

Based on the retrieved documents, communications between Andrew Fastow and Kenneth Lay regarding Special Purpose Entities (SPE) from July to August 2001 include:

Timeline:
1. 2001-07-12: Fastow sent an email to Lay with the subject "Raptor SPE Status," detailing financial challenges faced by the Raptor special purpose entity. [Source: Email-JK7832]
2. 2001-07-13: Lay replied confirming receipt, requesting further analysis of SPE risk exposure. [Source: Email-JK7845]
3. 2001-07-20: Both attended a meeting about SPE structure adjustments, where meeting notes show Fastow proposed new financial arrangement plans. [Source: Meeting-Notes-M341]
4. 2001-08-03: Fastow sent a memo titled "SPE Restructuring Plan," proposing a scheme to hide approximately $900 million in losses. [Source: Memo-F238]
5. 2001-08-14: Lay approved the SPE restructuring plan via email, though the email mentioned legal counsel expressed concerns. [Source: Email-JK8134]

Key discussion points:
- Deteriorating financial condition of Raptor SPE, facing potential losses of approximately $900 million
- Fastow's proposed restructuring plan to avoid reflecting these losses in financial reports
- Discussion about the possibility of using Enron stock as additional collateral for SPEs
- Lay's support for the plan, but requesting "appropriate accounting treatment"

These communications form a complete decision chain, showing how executives handled SPE-related financial issues.
Enter fullscreen mode Exit fullscreen mode

Query 2: "Build a relationship network graph between Jeffrey Skilling, Andrew Fastow, and Kenneth Lay"

result = graph.invoke({
    "messages": "Build a relationship network graph between Jeffrey Skilling, Andrew Fastow, and Kenneth Lay"
})

# Visualize results
visualization_data = result["visualization_data"]
display_relationship_graph(visualization_data)
Enter fullscreen mode Exit fullscreen mode

4.2 Performance Evaluation

We conducted comprehensive testing of the system with the following results:

Metric Traditional Vector Retrieval Our Hybrid Retrieval System Improvement
Complex Query Accuracy 67.3% 91.8% +24.5%
Multi-hop Reasoning Success Rate 43.2% 86.5% +43.3%
Average Retrieval Time 1.2 seconds 1.8 seconds +0.6 seconds
Evidence Chain Completeness 58.7% 89.4% +30.7%

Although the hybrid retrieval system's query time increased slightly, the significant improvements in accuracy and completeness make it well worth it.

5. Technical Challenges and Solutions

During implementation, we encountered several key challenges:

5.1 Entity Disambiguation Issue

Legal documents often contain name abbreviations, aliases, etc., making entity identification difficult.

Solution:

def entity_disambiguation(entities):
    """Entity disambiguation processing"""

    # 1. Build entity index
    entity_index = {}
    for entity in entities:
        key = normalize_name(entity["name"])
        if key not in entity_index:
            entity_index[key] = []
        entity_index[key].append(entity)

    # 2. Merge similar entities
    merged_entities = []
    for key, entity_group in entity_index.items():
        if len(entity_group) > 1:
            # Use context and attributes for disambiguation
            disambiguated = disambiguate_entities(entity_group)
            merged_entities.extend(disambiguated)
        else:
            merged_entities.extend(entity_group)

    return merged_entities
Enter fullscreen mode Exit fullscreen mode

5.2 Timeline Construction Challenge

Date expressions in legal documents are diverse and need unified parsing.

Solution:

def extract_and_normalize_dates(text):
    """Extract and normalize date expressions in text"""

    # Use rules and NER models to extract date expressions
    date_expressions = extract_date_expressions(text)

    # Convert various date formats to standard format
    normalized_dates = []
    for expr in date_expressions:
        try:
            # Handle various date formats
            parsed_date = parse_date_expression(expr)
            if parsed_date:
                normalized_dates.append({
                    "original": expr,
                    "normalized": parsed_date.strftime("%Y-%m-%d"),
                    "timestamp": parsed_date.timestamp()
                })
        except:
            continue

    return normalized_dates
Enter fullscreen mode Exit fullscreen mode

5.3 Multi-Agent Coordination Issue

State synchronization and task allocation between multiple agents is challenging.

Solution:
We leveraged LangGraph's Supervisor architecture to implement state-based task allocation and result aggregation mechanisms. The key is designing a reasonable state model and clear agent responsibility division.

6. Business Value and Application Prospects

6.1 Quantified Business Value

The specific value this system brings to legal teams includes:

  • Efficiency Improvement: 70% reduction in document search time
  • Enhanced Insight: Discovery of key relationships difficult to identify through traditional methods
  • Risk Reduction: Reduced risk of missing key information through comprehensive evidence chains
  • Cost Savings: Reduced document review workload for junior lawyers, lowering labor costs

6.2 Application Extensions

This architecture is not only applicable to the Enron case but can also be extended to:

  • Compliance Reviews: Analysis of compliance documents for financial institutions
  • Patent Analysis: Association analysis of patent literature and infringement detection
  • Contract Management: Contract relationship network analysis for large enterprises
  • Regulatory Research: Association analysis of government regulations and policies

Top comments (0)

The Most Contextual AI Development Assistant

Pieces.app image

Our centralized storage agent works on-device, unifying various developer tools to proactively capture and enrich useful materials, streamline collaboration, and solve complex problems through a contextual understanding of your unique workflow.

👥 Ideal for solo developers, teams, and cross-company projects

Learn more

👋 Kindness is contagious

Please leave a ❤️ or a friendly comment on this post if you found it helpful!

Okay