Building AI Agents The Better Way

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Data Pollution Audit Script

Audits a JSON file for records with data pollution using OpenAI's function calling.
Identifies records from other companies that have incorrectly crept into the target Account's data.

Usage:
    python audit.py <path_to_json_file>

Environment:
    OPENAI_API_KEY - Required for OpenAI API authentication
"""

import json
import os
import sys
from openai import OpenAI


def load_json_file(file_path):
    """Load and parse JSON file."""
    try:
        with open(file_path, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"Error: File not found: {file_path}")
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"Error: Invalid JSON format: {e}")
        sys.exit(1)


def extract_account_and_records(data):
    """Extract the Account object and all records from salesforce data."""
    try:
        records = data['salesforce']['records']
    except KeyError:
        print("Error: JSON structure must contain 'salesforce.records' array")
        sys.exit(1)

    # Find the Account object
    account = None
    for record in records:
        if record.get('Kind') == 'Account':
            account = record
            break

    if not account:
        print("Error: No Account object found in records (Kind: 'Account')")
        sys.exit(1)

    return account, records


def audit_with_openai(account, records):
    """Use OpenAI function calling to audit records for data pollution."""

    # Check for API key
    api_key = os.getenv('OPENAI_API_KEY')
    if not api_key:
        print("Error: OPENAI_API_KEY environment variable not set")
        sys.exit(1)

    # Initialize OpenAI client
    client = OpenAI(api_key=api_key)

    # Define the audit function schema
    audit_function = {
        "name": "audit_data_pollution",
        "description": "Report records from other companies that have incorrectly crept into the target Account's data",
        "parameters": {
            "type": "object",
            "properties": {
                "polluted_records": {
                    "type": "array",
                    "description": "List of records that belong to different companies/accounts and should not be associated with the target Account",
                    "items": {
                        "type": "object",
                        "properties": {
                            "record_index": {
                                "type": "integer",
                                "description": "The index of the polluted record in the records array"
                            },
                            "record": {
                                "type": "string",
                                "description": "Original record value as a JSON string."
                            },
                            "reason": {
                                "type": "string",
                                "description": "Explanation identifying which company this record actually belongs to and why it doesn't belong to the target Account"
                            }
                        },
                        "required": ["record_index", "record", "reason"]
                    }
                }
            },
            "required": ["polluted_records"],
            "additionalProperties": False
        }
    }

    # Prepare the prompt
    prompt = f"""You are a data quality auditor specializing in detecting company name pollution in Salesforce records.

TARGET ACCOUNT (the customer company you're auditing for):
{json.dumps(account, indent=2)}

ALL RECORDS TO AUDIT:
{json.dumps(records, indent=2)}

Your task - COMPANY NAME POLLUTION DETECTION:
Focus specifically on identifying records from OTHER COMPANIES that have incorrectly crept into this Account's data.

1. Examine each record for company names, account names, or organizational identifiers
2. Flag any record that clearly belongs to a DIFFERENT company/customer than the target Account
3. Look for company name mismatches in fields like: AccountName, Company, Organization, Customer, etc.
4. Identify records that reference different business entities or customer organizations

IMPORTANT EXCLUSIONS:
- Do NOT flag records simply because they mention "PubNub" or "pubnub" - PubNub is the service provider, not a competing company
- Do NOT flag minor data inconsistencies or formatting issues - focus ONLY on company name pollution
- Do NOT flag records that are legitimately related to the target Account (e.g., contacts, cases, opportunities for this company)

Call the audit_data_pollution function with your findings. For each polluted record, specify which OTHER COMPANY it actually belongs to.

If all records belong to the target Account company, return an empty polluted_records array.
"""

    # Make API call with forced function calling
    try:
        response = client.chat.completions.create(
            model="gpt-5",
            messages=[
                {
                    "role": "system",
                    "content": "You are a data quality auditor specializing in detecting company name pollution - identifying when records from other companies have incorrectly crept into a target customer's Salesforce data."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            tools=[{
                "type": "function",
                "function": audit_function
            }],
            tool_choice={
                "type": "function",
                "function": {"name": "audit_data_pollution"}
            }  # Force the function to be called
        )

        # Extract function call result
        tool_calls = response.choices[0].message.tool_calls
        if tool_calls and len(tool_calls) > 0:
            function_args = json.loads(tool_calls[0].function.arguments)
            return function_args.get('polluted_records', [])
        else:
            print("Warning: No function call was made by the API")
            return []

    except Exception as e:
        print(f"Error calling OpenAI API: {e}")
        sys.exit(1)


def display_results(polluted_records, total_records):
    """Display audit results."""
    print("\n" + "="*60)
    print("COMPANY NAME POLLUTION AUDIT REPORT")
    print("="*60)
    print(f"\nTotal records analyzed: {total_records}")
    print(f"Records from other companies found: {len(polluted_records)}")

    if polluted_records:
        print("\n" + "-"*60)
        print("RECORDS FROM OTHER COMPANIES:")
        print("-"*60)
        for item in polluted_records:
            record_idx = item.get('record_index')
            record = item.get('record')
            reason = item.get('reason')
            print(f"\n[Record #{record_idx}]")
            print(f"  Reason: {reason}")
            print(f"  Record: {record}")
        print("\n" + "="*60)
        print("STATUS: [WARNING] DATA POLLUTION DETECTED")
        print("="*60)
    else:
        print("\n" + "="*60)
        print("STATUS: [OK] NO DATA POLLUTION")
        print("="*60)


def main():
    """Main execution function."""
    # Check CLI arguments
    if len(sys.argv) < 2:
        print("Usage: python audit.py <path_to_json_file>")
        print("\nExample: python audit.py data.json")
        sys.exit(1)

    json_file = sys.argv[1]

    print(f"Loading JSON file: {json_file}")
    data = load_json_file(json_file)

    print("Extracting Account and records...")
    account, records = extract_account_and_records(data)
    print(f"Found target Account: {account.get('Name', 'N/A')}")
    print(f"Total records to audit: {len(records)}")

    print("\nAuditing for company name pollution with OpenAI...")
    polluted_records = audit_with_openai(account, records)

    display_results(polluted_records, len(records))


if __name__ == "__main__":
    main()