Untitled

mail@pastecode.io avatar
unknown
plain_text
16 days ago
5.1 kB
1
Indexable
Never
import json
from pydantic import BaseModel
from openai import OpenAI
from datetime import datetime

class UserInfo(BaseModel):
    subsidiary: str = ""
    brand: str = ""
    customer: str = ""
    start_period: str = ""
    end_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.

        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}]

    response = client.Completion.create(
        model="text-davinci-003",  # Adjust the model name as per your requirement
        prompt=json.dumps(messages),
        max_tokens=150  # Adjust max_tokens based on your needs
    )

    # Extract the response content
    completion = response.choices[0].text.strip()
    assistant_response = json.loads(completion)

    # Parse the assistant response into UserInfo
    user_info = UserInfo(**assistant_response)

    return user_info


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

    api_key = config['OPENAI_API_KEY']
    message = "Your input message here"

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

if __name__ == "__main__":
    main()
Leave a Comment