Untitled

import json
from pydantic import BaseModel
from openai import OpenAI

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 formatted as "YYYY-MM-DD to YYYY-MM-DD".
        """
    }

    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": "2023-01-01 to 2023-03-31"
            })
        },
        {
            "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": "2022-01-01 to 2022-12-31"
            })
        },
        {
            "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": "2024-03-01 to 2024-03-31"
            })
        },
        {
            "role": "user",
            "content": "List all overdue invoices for the brand DEF for Feb 2023."
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "",
                "brand": "DEF",
                "customer": "",
                "period": "2023-02-01 to 2023-02-28"
            })
        }
    ]

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

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

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

def format_period(period):
    period_map = {
        "Q1 2023": "2023-01-01 to 2023-03-31",
        "Q2 2023": "2023-04-01 to 2023-06-30",
        "Q3 2023": "2023-07-01 to 2023-09-30",
        "Q4 2023": "2023-10-01 to 2023-12-31",
        "Mar 2024": "2024-03-01 to 2024-03-31",
        "Feb 2023": "2023-02-01 to 2023-02-28",
        "2022": "2022-01-01 to 2022-12-31"
    }
    return period_map.get(period, 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 2022."

    user_info = get_user_info(api_key, message)
    user_info.period = format_period(user_info.period)
    print(user_info)

if __name__ == "__main__":
    main()