Untitled

import instructor
from pydantic import BaseModel
from openai import OpenAI
import json

# Define your desired output structure
class UserInfo(BaseModel):
    subsidiary: str = ""
    brand: str = ""
    customer: str = ""
    start_period: str = ""
    end_period: str = ""

def get_user_info(api_key, message):
    # Patch the OpenAI client
    client = instructor.from_openai(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.

        Figure out what is a subsidiary from the user prompt. A subsidiary is a company that is owned or controlled by another company, known as the parent or holding company. The parent company typically owns a majority stake in the subsidiary, allowing it to exert significant influence over its operations and strategic decisions. Subsidiaries operate as separate legal entities from their parent companies, which can be advantageous for legal and financial purposes, including taxation and liability management.
        Here the parent company is Doordash. So all the subsidiaries will be related to the Doordash parent company. If there is no subsidiary in the user prompt strictly return an empty string for subsidiary. 
        A customer refers to an individual, business, or organization that purchases goods or services from another party, typically a company or retailer. Customers play a crucial role in the economy by driving demand for products and services. Companies often segment customers based on various criteria such as demographics, preferences, and purchasing behavior to tailor their marketing strategies and improve customer satisfaction.
        Customers of Doordash could be businesses like Subway Inc or Chipotle Inc. If there is no customer in the user prompt strictly return an empty string for customer. If there is return the entire customer name. 

        A brand represents the identity and reputation of a company, product, or service in the marketplace. 
        An example of a brand is Subway. If there is no brand strictly return a empty string. 

        Some queries may not have brand, subsidiary or customer. return an empty string for them. 
        
        Special Instructions for Period: 
        Ensure that the start_period and end_period is normalized to the format "YYYY-MM-DD". Depending on the user prompt specify the start and end accordingly.  
        """
    }

    few_shot_examples = [
        {
            "role": "user",
            "content": "I need all transactions for Doordash Kitchens Inc for Q1 2023."
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "Doordash Kitchens Inc",
                "brand": "",
                "customer": "",
                "start_period": "2023-01-01",
                "end_period": "2023-03-31"
            })
        },
        {
            "role": "user",
            "content": "Show me the customer data for ID 12345 for 2022."
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "",
                "brand": "",
                "customer": "12345",
                "start_period": "2022-01-01",
                "end_period": "2022-12-31"
            })
        },
        {
            "role": "user",
            "content": "What is the total amount of receivables for the EMEA for Q2 2024?"
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "EMEA",
                "brand": "",
                "customer": "",
                "start_period": "2024-04-01",
                "end_period": "2024-06-30"
            })
        }, 
        {
            "role": "user",
            "content": "List all overdue invoices for the DEF for Feb 2023."
        },
        {
            "role": "assistant",
            "content": json.dumps({
                "subsidiary": "",
                "brand": "DEF",
                "customer": "",
                "start_period": "2023-02-01",
                "end_period": "2023-02-28"
            })
        }
    ]

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

    # Extract structured data from natural language
    user_info = client.chat.completions.create(
        model="gpt-4",
        response_model=UserInfo,
        messages=messages,
    )

    return user_info

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

    # Access configuration values
    api_key = config['OPENAI_API_KEY']
    # message = "Give me all bills & bill credits for the pre-paid rent account and subsidiary DDE-US for Q2?"
    message = "pre-paid rent account?"

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

if __name__ == "__main__":
    main()