Untitled

import json
from pydantic import BaseModel
from openai import OpenAI
from datetime import datetime

class UserInfo(BaseModel):
    subsidiary: str = ""
    brand: str = ""
    customer: str = ""
    period: str = ""

def get_user_info(api_key, message):
    client = OpenAI(api_key=api_key)

    system_prompt = {
        "role": "system",
        "content": """
        You are an assistant that extracts structured information from user prompts into specific entities: subsidiary, brand, customer, and period.
        Categorize the user query into one of the following:
        1. Subsidiary-related Queries: Require the entity "subsidiary"
        2. Customer-related Queries: Require the entity "customer"
        3. Brand-related Queries: Require the entity "brand"
        4. General Receivables Queries: Do not require specific entities (subsidiary, customer, brand)

        Based on the category, extract the necessary entities and leave others as empty strings. Ensure that the period is normalized to the format "MMM YYYY".
        """
    }

    few_shot_examples = [
        {
            "role": "user",
            "content": "I need all transactions for the subsidiary ABC Corp and the brand XYZ for the period Q1 2023."
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "ABC Corp",
                "brand": "XYZ",
                "customer": "",
                "period": "Jan 2023 - Mar 2023"
            })
        },
        {
            "role": "user",
            "content": "Show me the customer data for customer ID 12345 for the period 2022."
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "",
                "brand": "",
                "customer": "12345",
                "period": "Jan 2022 - Dec 2022"
            })
        },
        {
            "role": "user",
            "content": "What is the total amount of receivables for the subsidiary EMEA for Mar 2024?"
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "EMEA",
                "brand": "",
                "customer": "",
                "period": "Mar 2024"
            })
        },
        {
            "role": "user",
            "content": "List all overdue invoices for the brand DEF for Feb 2023."
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "",
                "brand": "DEF",
                "customer": "",
                "period": "Feb 2023"
            })
        }
    ]

    messages = [system_prompt] + few_shot_examples + [{"role": "user", "content": message}]

    response = client.chat_completions.create(
        model="gpt-4",
        messages=messages
    )

    user_info = UserInfo.parse_raw(response.choices[0].message["content"])

    user_info.period = normalize_period(user_info.period)

    return user_info

def normalize_period(period):
    current_date = datetime.now()

    period_map = {
        "current month": current_date.strftime("%b %Y"),
        "current quarter": f"Q{((current_date.month - 1) // 3) + 1} {current_date.year}",
        "current year": f"{current_date.year}",
        "fiscal year": f"{current_date.year}"
    }

    normalized_period = period_map.get(period.lower(), period)
    
    # Additional normalization logic for common date formats
    if period.lower() in period_map:
        return normalized_period
    try:
        # Attempt to parse common date formats
        normalized_period = datetime.strptime(period, "%b %Y").strftime("%b %Y")
    except ValueError:
        pass

    return normalized_period

def main():
    with open('config.json', 'r') as f:
        config = json.load(f)

    api_key = config['OPENAI_API_KEY']
    message = "Show me the customer data for customer ID 12345 for the period current year."

    user_info = get_user_info(api_key, message)
    print(user_info)

if __name__ == "__main__":
    main()