app.py

import os
import re
import json
import gradio as gr
from openai import OpenAI
import gspread
from google.oauth2.service_account import Credentials

SCOPES = [
    "https://www.googleapis.com/auth/spreadsheets",
    "https://www.googleapis.com/auth/drive"
]

# Initialize the OpenAI client with the API key from environment variables.
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

# In-memory set to track submitted emails (this resets when the app restarts).
submitted_emails = set()

def get_google_sheet():
    """
    Connects to the Google Sheet using service account credentials stored
    in the environment variable "GOOGLE_CREDS_JSON" and returns the worksheet
    named "Submissions" from the spreadsheet identified by "SPREADSHEET_ID".
    """
    creds = Credentials.from_service_account_info(
        json.loads(os.environ["GOOGLE_CREDS_JSON"]),
        scopes=SCOPES
    )
    gc = gspread.authorize(creds)
    sh = gc.open_by_key(os.environ["SPREADSHEET_ID"])
    worksheet = sh.worksheet("Submissions")
    return worksheet

def get_evaluation_questions():
    """
    Loads evaluation questions and expected answers from environment variables.
    
    Expected environment variables:
    - TEST_QUESTION_1: a JSON array of user query strings.
    - TEST_EXPECTED_1: a JSON array of JSON-like strings representing expected outputs.
    
    Both lists must be of equal length.
    """
    questions_str = os.environ.get("TEST_QUESTION_1")
    docs_str = os.environ.get("TEST_DOCUMENTS_1")
    expected_str = os.environ.get("TEST_EXPECTED_1")

    if not questions_str or not expected_str or not docs_str:
        return []
    try:
        questions_list = json.loads(questions_str)
    except Exception as e:
        print(f"Error parsing questions: {str(e)}")
        return []
    try:
        expected_list = json.loads(expected_str)
    except Exception as e:
        print(f"Error parsing expected answers: {str(e)}")
        return []
    try:
        docs_list = json.loads(docs_str)
    except Exception as e:
        print(f"Error parsing documents: {str(e)}")
        return []
    
    # Ensure all lists are of the same length.
    if len(questions_list) != len(expected_list) or len(questions_list) != len(docs_list):
        print("Mismatch in length: questions list and expected answers list must have the same length.")
        return []
    
    return [{"question": q, "expected": e, "docs": d} for q, e, d in zip(questions_list, expected_list, docs_list)]

# Load evaluation questions at startup.
EVALUATION_QUESTIONS = get_evaluation_questions()

def sanitize_input(text):
    """
    Sanitizes input to allow only alphanumerics and some punctuation,
    then truncates to 500 characters.
    """
    clean_text = re.sub(r"[^a-zA-Z0-9\s.,!?@:\-]", "", text)
    return clean_text.strip()[:500]

def sanitize_prompt(text):
    """
    Sanitizes the system prompt by stripping and limiting its length.
    """
    return text.strip()[:8000]

def validate_email(email):
    """
    Validates that the provided email is in a valid format.
    Returns True if valid, False otherwise.
    """
    email_regex = r"^[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+$"
    return re.match(email_regex, email) is not None

def submit_prompt(email, name, system_prompt):
    """
    Handles the full submission process:
     - Validates email format.
     - Checks if the email has already been used (by in-memory set and Google Sheet).
     - Sanitizes input fields.
     - Processes the system prompt against each evaluation question using the OpenAI API.
     - For each test question, records the verdict and answer.
     - Appends the submission as a new row in the Google Sheet with columns:
         Name, Email, System Prompt, Score, and for each of the 7 test questions: verdict and answer.
    Returns a result message with evaluation details.
    """
    # Validate email format.
    if not validate_email(email):
        return "Invalid email address. Please enter a valid email."
    
    # Check if this email has already been submitted (in-memory).
    if email in submitted_emails:
        return f"Submission already received for {email}. You can only submit once."
    
    # Connect to Google Sheet and check if the email already exists.
    try:
        sheet = get_google_sheet()
        email_col = sheet.col_values(2)  # Assumes column 2 contains the email addresses.
        if email in email_col[1:]:  # Skip header row.
            return f"Submission already received for {email}. You can only submit once."
    except Exception as e:
        print(f"Error accessing Google Sheet: {str(e)}")
        return f"Error accessing Google Sheet: {str(e)}"
    
    # Sanitize inputs.
    email = sanitize_input(email)
    name = sanitize_input(name)
    system_prompt = sanitize_prompt(system_prompt)
    
    score = 0
    responses = []    # For display output.
    verdicts = []     # For storing each question's verdict in the sheet.
    answers_list = [] # For storing each question's answer in the sheet.
    
    # Process each evaluation question.
    for item in EVALUATION_QUESTIONS:
        question = item["question"]
        docs = item["docs"]
        expected = item["expected"]
        try:
            response = client.chat.completions.create(
                model="gpt-4o-mini",  # Ensure this model identifier matches your deployed model.
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": question}
                ]
            )
            answer = response.choices[0].message.content.strip()
        except Exception as e:
            answer = f"Error during OpenAI API call: {str(e)}"
        
        verdict = ""
        # Check if the answer is a valid JSON.
        try:
            parsed_answer = json.loads(answer)
            answer_to_store = json.dumps(parsed_answer)  # Normalize parsed JSON as string.
        except json.JSONDecodeError as e:
            verdict = f"Incorrect (Invalid JSON: {str(e)})"
            responses.append(
                f"Question: {question}\n"
                f"Answer: {answer}\n"
                f"Expected: {json.dumps(expected)}\n"
                f"Result: {verdict}\n"
            )
            verdicts.append(verdict)
            answers_list.append(answer)
            continue
        
        # Verify that all required keys are present.
        required_keys = ["document_level", "clause_level"]
        missing_keys = [key for key in required_keys if key not in parsed_answer]
        if missing_keys:
            verdict = f"Incorrect (Missing Keys: {', '.join(missing_keys)})"
            responses.append(
                f"Question: {question}\n"
                f"Answer: {json.dumps(parsed_answer)}\n"
                f"Expected: {json.dumps(expected)}\n"
                f"Result: {verdict}\n"
            )
            verdicts.append(verdict)
            answers_list.append(json.dumps(parsed_answer))
            continue
        
        # Compare values for each required key.
        incorrect_values = []
        for key in required_keys:
            if parsed_answer[key] != expected[key]:
                incorrect_values.append(key)
        
        if len(incorrect_values) == 2:
            verdict = "Incorrect (Both values are incorrect)"
        elif len(incorrect_values) == 1:
            verdict = f"Incorrect (Value for key '{incorrect_values[0]}' is incorrect)"
        else:
            score += 1
            verdict = "Correct"
        
        responses.append(
            f"Question: {question}\n"
            f"Answer: {json.dumps(parsed_answer)}\n"
            f"Expected: {json.dumps(expected)}\n"
            f"Result: {verdict}\n"
        )
        verdicts.append(verdict)
        answers_list.append(json.dumps(parsed_answer))
    
    result_details = "\n".join(responses)
    
    # Record this email locally so that subsequent submissions are blocked.
    submitted_emails.add(email)
    
    # Prepare the row for Google Sheets:
    # The row format is: Name, Email, System Prompt, Score, then for each of the 7 test questions: Verdict, Answer.
    row = [name, email, system_prompt, str(score)]
    for v, a in zip(verdicts, answers_list):
        row.extend([v, a])
    
    # Append the new row to the Google Sheet.
    try:
        sheet.append_row(row)
    except Exception as e:
        print(f"Error appending row to Google Sheet: {str(e)}")
        return f"Error saving submission: {str(e)}"
    
    return (
        f"Thank you for your submission, {name}!\n\n"
    )

def build_interface():
    """
    Constructs the Gradio interface with a submission button and single-submission mechanism.
    """
    with gr.Blocks() as demo:
        gr.Markdown("# Applicant Task: Target Company & Law Firm Identification")
        gr.Markdown("## Identifying Parties, Law Firms, and Target Company Presence")
        
        # General description
        gr.Markdown("""
        This task involves processing a user query to determine the relevance to the intended task, followed by analyzing textual data to extract information about law firms representing parties (Buyer, Seller, and Third Parties) and verifying the presence of a target company. 

        The system is designed to sequentially leverage three LLM functions:
        
        **Step 1:** LLM1 determines if the user's query mentions any target company.  
        - If no target company is found, LLM1 responds with a message wrapped in `<user_message></user_message>` XML tags to inform the user that the query is irrelevant to this task.  
        - If the query contains a target company, LLM1 moves forward with a formatted acknowledgment of the identified target company.

        **Step 2:** LLM2 examines four separate paragraphs provided as input. For each paragraph, it extracts specific information about:
        - The Buyer's representative law firm.
        - The Seller's representative law firm.
        - Any third-party law firm present.
        - Whether the target company is mentioned in the paragraph.
        Each paragraph's results are formatted and concatenated for the next step.

        **Step 3:** LLM3 compiles the information from all analyzed paragraphs and outputs a structured JSON object with the following keys:
        
        ```json
        {
        "buyer_firm": "string",
        "seller_firm": "string",
        "third_party": "string",
        "contains_target_firm": boolean
        }
        ```
        
        The goal is to identify the representative law firms of involved parties and determine if the target company is mentioned, ensuring the results are structured and accurate.
        
        **Key Considerations:**
        - The output must adhere to the prescribed JSON format for the final step.
        - Ensure the system can accurately extract and classify relevant information from the input paragraphs.
        """)

        # Example Inputs and Outputs in an Accordion
        with gr.Accordion("Example Workflow", open=False):
            gr.Markdown("""
            **Example Query and System Output:**

            **User Query:**  
            *"Is Kirkland present?"*

            Step 1 (LLM1):  
            - If no target company is identified:  
            Output: `<user_message>Query is not relevant to the intended task.</user_message>`  

            - If a target company is identified:  
            Output: *"The query mentions the target company Kirkland."*

            Step 2 (LLM2 for Paragraphs):  
            **Input Paragraph Example:**  
            *"Representation agreements between Buyers and Kirkland & Ellis are included."*  

            **Output:**  
            *"Buyer Firm: Kirkland & Ellis, Seller Firm: None, Third Party Firm: None, Contains Target Firm: True."*  

            Step 3 (LLM3 Final Output):  
            Compiled JSON:
            ```json
            {
            "buyer_firm": "Kirkland & Ellis",
            "seller_firm": null,
            "third_party": null,
            "contains_target_firm": true
            }
            ```
            """)

        # Challenge instructions and testing guidance
        with gr.Accordion("Task Instructions and Testing", open=False):
            gr.Markdown("""
            - Design prompts that ensure proper interaction between the three LLM systems, with each step contributing to the final output.
            - Ensure strict adherence to JSON formatting requirements (e.g., no extra characters that may cause JSON parsing errors).
            - Test extensively to verify accurate law firm and target company identification.
            
            **Output Requirements:**
            - Ensure final LLM3 JSON output has the following keys:
            - `"buyer_firm"`
            - `"seller_firm"`
            - `"third_party"`
            - `"contains_target_firm"`
            - Values must be accurately extracted or classified based on LLM2's parsed data.

            **Hints for Crafting System Prompts:**
            - Explicitly specify formatting requirements at each step.
            - Clarify the task definitions and expected classifications in each system prompt for LLM1, LLM2, and LLM3.
            - Test using diverse sample data for robustness.

            You can only submit once, so validate your system prompts thoroughly using mock queries and example data before final submission.
            
            Good Luck!
            """)

        gr.Markdown("""
        Enter your name and email below, as listed in your CV, and submit your designed prompts.
        
        Remember: Focus on clarity, accuracy, and structured responses to achieve a high score!
        """)


        email_input = gr.Textbox(label="Email", placeholder="[email protected]")
        name_input = gr.Textbox(label="First Name, Last Name", placeholder="John, Smith")
        system_prompt_input = gr.Textbox(
            label="System Prompt",
            placeholder="Enter your system prompt here...",
            lines=6,
        )
        submit_button = gr.Button("Submit")
        output_text = gr.Textbox(label="Results", lines=15)

        submit_button.click(
            fn=submit_prompt,
            inputs=[email_input, name_input, system_prompt_input],
            outputs=output_text,
        )
        
        
    
    return demo

if __name__ == "__main__":
    interface = build_interface()
    # Launch the app on 0.0.0.0 so it is accessible externally (e.g., in a container).
    interface.launch(server_name="0.0.0.0", server_port=7860)