AI code gen for report building using Claude Sonnet 3.5

report_code_gen.py

import argparse
import os
import anthropic
from pymongo import MongoClient

anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')

client = anthropic.Anthropic(
    api_key=anthropic_api_key,
)

model = "claude-3-5-sonnet-20241022"
parser = argparse.ArgumentParser(description='Process a report description and generate a report.')
parser.add_argument('--input', type=str, required=True, help='The report description (brand, fields, etc.)')
args = parser.parse_args()    
mongocli = MongoClient(mongo_url=os.getenv('MONGO_URI'), username=os.getenv('MONGO_USERNAME'), password=os.getenv('MONGO_PASSWORD'))
db = mongocli.blintz

company_codes = [company['code'] for company in db.companies.find({}, {'code': 1})]

COLLECTIONS = [
    {
        'name': 'users', 
        'description': 'System users, i.e. real estate agents. If you need to use an "agents" collection, this is the one.',
        'fields': {
            '_id': 'The ObjectId is the MongoDB id that uniquely identifies a user document.', 
            'name': 'A string representing the user name', 
            'email': 'A string representing the user email address', 
            'phone': 'A string representing the user phone number',
            'company_id': "A string that identifies the company the agent belongs to. Important: this is not the company code.",
        }
    },
    {
        'name': 'deals', 
        'description': 'Deals are transaction opportunities between a lead and an owner. They are managed by agents.',
        'fields': {
            '_id': 'The ObjectId is the MongoDB id that uniquely identifies a deal document.', 
            'company_id': 'The ObjectId of the company that owns the deal. Important: this is not the company code.', 
            'agent_id': 'The ObjectId of the agent managing the deal. Its string representation is {"agent_id: {"$oid": "the id"}}', 
            'reference': 'A string representing the listing commercial identifier that the deal is related to', 
            'origin': 'A string representing the origin of the deal, such as instagram, real estate portal, Pilar Homes, etc.',
            'created_at': 'A string representing the date and time the deal was created in format YYYY-mm-dd HH:MM:SS',
            'transaction_type': 'A string that represents the type of transaction, such as sale, rent, sale/rent'
        },
    },
    {
        'name': 'companies',
        'description': 'Companies are the real estate brokerages. If the user provides a code to filter the data, it will be a company code. The _id may be retrieved by querying the company with the given code. Company codes are not used to join data.',
        'fields': {
            '_id': 'The ObjectId is the MongoDB id that uniquely identifies a company document. Its string representation is \"{"$oid": "the id"}\"',
            'code': 'The company code is a short and human friendly string that uniquely identifies the company. Never use it for joining data.',
            'name': 'A string representing the company name',
        }
    },
    {
        'name': 'properties',
        'description': 'Properties are real estate listings. They are managed by agents (users) and owned by an owner (contact). They belong to a company.',
        'fields': {
            '_id': 'The ObjectId is the MongoDB id that uniquely identifies a property document.',
            'company_id': 'A string representing the id of the company that registered the listing. Important: this is not an ObjectId. Use it as \"{\"company_id\": \"61f7b7b7b7b7b7b7b7b7b7b7\"}\". Important: this is not the company code.',
            'agent_id': 'A string representing the id of the agent managing the property. Its string representation is {"agent_id": {"$oid": "the id"}}',
            'commercial_id': 'The commercial id uniquely identifies a listing. It is a short and human friendly string. Never use it for joining data.',
            'address': 'A string representing the property address',
            'city': 'A string representing the property city',
            'state': 'A string representing the property state',
            'country': 'A string representing the property country',
            'asking_price': 'A number representing the property asking price',
            'deactivation_details': 'A string representing the reason why the property was deactivated. If it is null or empty, the property is active.',
            'people': 'An array of dictionaries representing the people associated with the property. Each dictionary has a "role" key with the person role and a "contact_id" key with the person contact id. The role may be "owner" or "lead".',
        }
    },
    {
        'name': 'contacts',
        'description': 'Contacts may be owners of properties or leads interested in purchasing or renting a property. They belong to a company.',
        'fields': {
            '_id': 'The ObjectId is the MongoDB id that uniquely identifies a contact document.',
            'company_id': 'The ObjectId is the MongoDB id that uniquely identifies the company that created the contact. Important: this is not the company code.',
            'name': 'A string representing the contact name',
            'emails': 'An array of strings representing the contact email addresses',
            'phones': 'An array of dictionaries representing the contact phone numbers. Each dictionary has a "number" key with the phone number and a "country_code" key with the country code',
        }
    }
]

def generate_report_building_code(report_description: str) -> str:
    """Generate a Python script for building the report requested by the user.

    Args:
        report_description (str): the report description

    Returns:
        str: the Python code that builds a report
    """
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": f"{report_description}.\nAvailable collections: {COLLECTIONS}\nCompany codes: {company_codes}\n.Always demand a company code from the user to filter the data -- the user may use the terms imobiliária, marca, brand or company to reference a company. If the user wants a field that does not exist in a collection, don't add it to the report and don't ask the user for the field.",
                }
            ]
        }
    ]
    response = client.messages.create(
        model=model,
        max_tokens=1000,
        temperature=0,
        system="Generate Python code to create a CSV report given a natural language input. Insert print statements for debugging and following the script execution progress. Just output the script. If any CSV field is null, use 'N/A', do not skip fields unless the user explicitly asks for it. Use projections to improve query performance.",
        messages=messages
    )
    
    return response.content[0].text

if __name__ == '__main__':
    report_description = args.input
    
    report_code = generate_report_building_code(report_description)
    
    print(report_code)