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
    period: str

def get_user_info(api_key, message):
    # Patch the OpenAI client
    client = instructor.from_openai(OpenAI(api_key=api_key))

    # Define the system prompt and few-shot examples
    system_prompt = {
        "role": "system",
        "content": "You are an assistant that extracts structured information from user prompts."
    }

    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": "Q1 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": "2022"
            })
        }
    ]

    # Prepare the messages list with system prompt, few-shot examples, and the user message
    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)

    # Convert the Pydantic object to JSON
    # user_info_json = json.dumps(user_info.dict())

    # print(user_info_json)

if __name__ == "__main__":
    main()