# install packages
!pip install anthropic pypdf pandas matplotlib numpy rouge-score nltk seaborn --quiet

import re
 
import anthropic
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
 
# Set up Anthropic client
# You can set up a .env file with your API key to keep it private, and import it like so:
# from dotenv import load_dotenv
# load_dotenv()
 
# or add your key directly
api_key = "ANTHROPIC_API_KEY"  # Replace ANTHROPIC_API_KEY with your actual API key
client = anthropic.Anthropic(api_key=api_key)
 
print("Setup complete!")

import re
 
import pypdf
 
pdf_path = "data/Sample Sublease Agreement.pdf"
 
 
def extract_text_from_pdf(pdf_path):
    with open(pdf_path, "rb") as file:
        reader = pypdf.PdfReader(file)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + "\n"
    return text
 
 
def clean_text(text):
    # Remove extra whitespace
    text = re.sub(r"\s+", " ", text)
    # Remove page numbers
    text = re.sub(r"\n\s*\d+\s*\n", "\n", text)
    return text.strip()
 
 
def prepare_for_llm(text, max_tokens=180000):
    # Truncate text to fit within token limit (approximate)
    return text[: max_tokens * 4]  # Assuming average of 4 characters per token
 
 
def get_llm_text(path):
    extracted_text = extract_text_from_pdf(path)
    cleaned_text = clean_text(extracted_text)
    llm_ready_text = prepare_for_llm(cleaned_text)
    return llm_ready_text
 
 
# You can now use get_llm_text in your LLM prompt
text = get_llm_text(pdf_path)
print(text[:500])

def basic_summarize(text, max_tokens=1000):
    # Prompt the model to summarize the text
    prompt = f"""Summarize the following text in bullet points. Focus on the main ideas and key details:
    {text}
    """
 
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=max_tokens,
        system="You are a legal analyst known for highly accurate and detailed summaries of legal documents.",
        messages=[
            {"role": "user", "content": prompt},
            {
                "role": "assistant",
                "content": "Here is the summary of the legal document: <summary>",
            },
        ],
        stop_sequences=["</summary>"],
    )
 
    return response.content[0].text

basic_response = basic_summarize(text, max_tokens=1000)
 
print(basic_response)

# We import from our data directory to save space in our notebook
from data.multiple_subleases import document1, document2, document3, sample1, sample2, sample3
 
 
def basic_summarize_multishot(text, max_tokens=1000):
    # Prompt the model to summarize the text
    prompt = f"""Summarize the following text in bullet points. Focus on the main ideas and key details:
        {text}
 
    Do not preamble.
 
    Use these examples for guidance in summarizing:
 
    <example1>
        <original1>
            {document1}
        </original1>
 
        <summary1>
            {sample1}
        </summary1>
    </example1>
 
    <example2>
        <original2>
            {document2}
        </original2>
 
        <summary2>
            {sample2}
        </summary2>
    </example2>
 
    <example3>
        <original3>
            {document3}
        </original3>
 
        <summary3>
            {sample3}
        </summary3>
    </example3>
    """
 
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=max_tokens,
        system="You are a legal analyst known for highly accurate and detailed summaries of legal documents.",
        messages=[
            {"role": "user", "content": prompt},
            {
                "role": "assistant",
                "content": "Here is the summary of the legal document: <summary>",
            },
        ],
        stop_sequences=["</summary>"],
    )
 
    return response.content[0].text

basic_multishot_response = basic_summarize_multishot(text, max_tokens=1000)
 
print(basic_multishot_response)

def guided_legal_summary(text, max_tokens=1000):
    # Prompt the model to summarize the text
    prompt = f"""Summarize the following legal document. Focus on these key aspects:
 
    1. Parties involved
    2. Main subject matter
    3. Key terms and conditions
    4. Important dates or deadlines
    5. Any unusual or notable clauses
 
    Provide the summary in bullet points under each category.
 
    Document text:
    {text}
 
    """
 
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=max_tokens,
        system="You are a legal analyst known for highly accurate and detailed summaries of legal documents.",
        messages=[
            {"role": "user", "content": prompt},
            {
                "role": "assistant",
                "content": "Here is the summary of the legal document: <summary>",
            },
        ],
        stop_sequences=["</summary>"],
    )
 
    return response.content[0].text

# Example usage
legal_summary = guided_legal_summary(text)
print(legal_summary)

def guided_sublease_summary(text, model="claude-sonnet-4-5", max_tokens=1000):
    # Prompt the model to summarize the sublease agreement
    prompt = f"""Summarize the following sublease agreement. Focus on these key aspects:
 
    1. Parties involved (sublessor, sublessee, original lessor)
    2. Property details (address, description, permitted use)
    3. Term and rent (start date, end date, monthly rent, security deposit)
    4. Responsibilities (utilities, maintenance, repairs)
    5. Consent and notices (landlord's consent, notice requirements)
    6. Special provisions (furniture, parking, subletting restrictions)
 
    Provide the summary in bullet points nested within the XML header for each section. For example:
 
    <parties involved>
    - Sublessor: [Name]
    // Add more details as needed
    </parties involved>
 
    If any information is not explicitly stated in the document, note it as "Not specified". Do not preamble.
 
    Sublease agreement text:
    {text}
 
    """
 
    response = client.messages.create(
        model=model,
        max_tokens=max_tokens,
        system="You are a legal analyst specializing in real estate law, known for highly accurate and detailed summaries of sublease agreements.",
        messages=[
            {"role": "user", "content": prompt},
            {
                "role": "assistant",
                "content": "Here is the summary of the sublease agreement: <summary>",
            },
        ],
        stop_sequences=["</summary>"],
    )
 
    return response.content[0].text

# Example usage
sublease_summary = guided_sublease_summary(text)
print(sublease_summary)

import re
 
 
def parse_sections_regex(text):
    pattern = r"<(.*?)>(.*?)</\1>"
    matches = re.findall(pattern, text, re.DOTALL)
 
    parsed_sections = {}
    for tag, content in matches:
        items = [item.strip("- ").strip() for item in content.strip().split("\n") if item.strip()]
        parsed_sections[tag] = items
 
    return parsed_sections
 
 
# Parse the sections
parsed_sections = parse_sections_regex(sublease_summary)
 
# Check if parsing was successful
if isinstance(parsed_sections, dict) and "parties involved" in parsed_sections:
    print("Parties involved:")
    for item in parsed_sections["parties involved"]:
        print(f"- {item}")
else:
    print("Error: Parsing failed or 'parties involved' section not found.")
    print("Parsed result:", parsed_sections)

from data.multiple_subleases import document1, document2, document3
 
 
def chunk_text(text, chunk_size=2000):
    return [text[i : i + chunk_size] for i in range(0, len(text), chunk_size)]
 
 
def summarize_long_document(text, max_tokens=2000):
    chunks = chunk_text(text)
 
    # Iterate over chunks and summarize each one
    # We use guided_legal_summary here, but you can use basic_summarize or any other summarization function
    # Note that we'll also use haiku for the interim summaries, and the 3.5 sonnet for the final summary
    chunk_summaries = [
        guided_sublease_summary(chunk, model="claude-haiku-4-5", max_tokens=max_tokens)
        for chunk in chunks
    ]
 
    final_summary_prompt = f"""
 
    You are looking at the chunked summaries of multiple documents that are all related. Combine the following summaries of the document from different truthful sources into a coherent overall summary:
 
    {"".join(chunk_summaries)}
 
    1. Parties involved (sublessor, sublessee, original lessor)
    2. Property details (address, description, permitted use)
    3. Term and rent (start date, end date, monthly rent, security deposit)
    4. Responsibilities (utilities, maintenance, repairs)
    5. Consent and notices (landlord's consent, notice requirements)
    6. Special provisions (furniture, parking, subletting restrictions)
 
    Provide the summary in bullet points nested within the XML header for each section. For example:
 
    <parties involved>
    - Sublessor: [Name]
    // Add more details as needed
    </parties involved>
 
    If any information is not explicitly stated in the document, note it as "Not specified".
 
    Summary:
    """
 
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=max_tokens,
        system="You are a legal expert that summarizes notes on one document.",
        messages=[
            {"role": "user", "content": final_summary_prompt},
            {
                "role": "assistant",
                "content": "Here is the summary of the legal document: <summary>",
            },
        ],
        stop_sequences=["</summary>"],
    )
 
    return response.content[0].text
 
 
# Example usage
# combine 3 documents (all related) together
text = document1 + document2 + document3
long_summary = summarize_long_document(text)
print(long_summary)

class LegalSummaryIndexedDocuments:
    def __init__(self, client):
        self.client = client  # Claude client
        self.documents: list[dict[str, str]] = []  # List of docs to store
        self.summaries: list[str] = []
 
    def add_document(self, doc_id: str, content: str):
        # Adds a document to the index
        self.documents.append({"id": doc_id, "content": content})
 
    def generate_summaries(self):
        # Generates summaries for all documents in the index
        for doc in self.documents:
            summary = self._generate_legal_summary(doc["content"])
            self.summaries.append(summary)
 
    def _generate_legal_summary(self, content: str) -> str:
        # Note how we constrain the content to a maximum of 2000 words. We do this because we don't need that much information for the intial ranking.
        prompt = f"""
        Summarize the following sublease agreement. Focus on these key aspects:
 
        1. Parties involved (sublessor, sublessee, original lessor)
        2. Property details (address, description, permitted use)
        3. Term and rent (start date, end date, monthly rent, security deposit)
        4. Responsibilities (utilities, maintenance, repairs)
        5. Consent and notices (landlord's consent, notice requirements)
        6. Special provisions (furniture, parking, subletting restrictions)
 
        Provide the summary in bullet points nested within the XML header for each section. For example:
 
        <parties involved>
        - Sublessor: [Name]
        // Add more details as needed
        </parties involved>
 
        If any information is not explicitly stated in the document, note it as "Not specified".
 
        Sublease agreement text:
        {content[:2000]}...
 
        Summary:
        """
 
        response = client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=500,
            temperature=0.2,
            messages=[
                {"role": "user", "content": prompt},
                {
                    "role": "assistant",
                    "content": "Here is the summary of the legal document: <summary>",
                },
            ],
            stop_sequences=["</summary>"],
        )
        return response.content[0].text
 
    def rank_documents(self, query: str, top_k: int = 3) -> list[tuple[str, float]]:
        """
        Rank documents based on their relevance to the given query.
        We use Haiku here as a cheaper, faster model for ranking.
        """
        ranked_scores = []
        for summary in self.summaries:
            prompt = f"Legal document summary: {summary}\n\nLegal query: {query}\n\nRate the relevance of this legal document to the query on a scale of 0 to 10. Only output the numeric value:"
 
            response = client.messages.create(
                model="claude-haiku-4-5",
                max_tokens=2,
                temperature=0,
                messages=[{"role": "user", "content": prompt}],
            )
            ranked_score = float(response.content[0].text)
            ranked_scores.append(ranked_score)
 
        ranked_indices = np.argsort(ranked_scores)[::-1][:top_k]
        return [(self.documents[i]["id"], ranked_scores[i]) for i in ranked_indices]
 
    def extract_relevant_clauses(self, doc_id: str, query: str) -> list[str]:
        """
        Extracts relevant clauses from a document based on a query.
        """
        doc_content = next(doc["content"] for doc in self.documents if doc["id"] == doc_id)
 
        prompt = f"""
        Given the following legal query and document content, extract the most relevant clauses or sections and write the answer to the query.
        Provide each relevant clause or section separately, preserving the original legal language:
 
        Legal query: {query}
 
        Document content: {doc_content}...
 
        Relevant clauses or sections (separated by '---'):"""
 
        response = client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=1000,
            temperature=0,
            messages=[{"role": "user", "content": prompt}],
        )
 
        clauses = re.split(r"\n\s*---\s*\n", response.content[0].text.strip())
        return [clause.strip() for clause in clauses if clause.strip()]

from data.multiple_subleases import document1, document2, document3
 
lsid = LegalSummaryIndexedDocuments(client=client)
 
# Add documents
lsid.add_document("doc1", document1)
lsid.add_document("doc2", document2)
lsid.add_document("doc3", document3)
 
# Generate summaries - this would happen at ingestion
lsid.generate_summaries()
 
# Rank documents for a legal query
legal_query = "What contract is for the sublessor Apex Innovations, LLC?"
ranked_results = lsid.rank_documents(legal_query)
 
print("Initial ranking:", ranked_results)
 
# Extract relevant clauses from the top-ranked document
top_doc_id = ranked_results[0][0]
relevant_clauses = lsid.extract_relevant_clauses(top_doc_id, legal_query)
 
print("\nRelevant clauses from the top-ranked document:")
for i, clause in enumerate(relevant_clauses[1:], 1):
    print(f"Clause {i}: {clause}")

import re
 
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
 
%matplotlib inline
plt.style.use("seaborn")
 
# Load the data
df = pd.read_csv("data/results.csv")
 
 
# Function to extract PASS/FAIL and score
def extract_result(text):
    match = re.search(r"\[(PASS|FAIL)\]\s*\((\d+\.\d+)\)", str(text))
    if match:
        return match.group(1), float(match.group(2))
    return "UNKNOWN", 0.0
 
 
# Apply the extraction to relevant columns
for col in df.columns[2:]:
    df[f"{col}_result"], df[f"{col}_score"] = zip(*df[col].apply(extract_result), strict=False)
 
# Prepare data for grouped accuracy score
models = ["3.5 Sonnet", "3.0 Haiku"]
prompts = ["basic_summarize", "guided_legal_summary", "summarize_long_document"]
 
results = []
for model in models:
    for prompt in prompts:
        col = f"[{model}] prompts.py:{prompt}_result"
        if col in df.columns:
            pass_rate = (df[col] == "PASS").mean()
            results.append({"Model": model, "Prompt": prompt, "Pass Rate": pass_rate})
 
result_df = pd.DataFrame(results)
 
# 1. Grouped bar chart for accuracy scores
plt.figure(figsize=(12, 6))
result_pivot = result_df.pivot(index="Prompt", columns="Model", values="Pass Rate")
result_pivot.plot(kind="bar")
plt.title("Pass Rate by Model and Prompt")
plt.ylabel("Pass Rate")
plt.legend(title="Model")
plt.xticks(rotation=45, ha="right")
plt.tight_layout()
plt.show()
 
# 2. Box plot of scores
plt.figure(figsize=(8, 8))
score_cols = [col for col in df.columns if col.endswith("_score")]
score_data = df[score_cols].melt()
sns.boxplot(x="variable", y="value", data=score_data)
plt.title("Distribution of Scores")
plt.xticks(rotation=90)
plt.xlabel("Model and Prompt")
plt.ylabel("Score")
plt.tight_layout()
plt.show()
 
# Display summary statistics
summary_stats = df[[col for col in df.columns if col.endswith("_score")]].describe()
display(summary_stats)