Untitled

import json
from concurrent.futures import ThreadPoolExecutor
from pydantic_user_prompt import get_user_info
from subsidiary_search_client import UserPrompt as SubsidiaryUserPrompt, handle_user_prompt_for_Subsidiary, format_subsidiary_search_results
from customer_search_client import UserPrompt as CustomerUserPrompt, handle_user_prompt_for_Customer, format_customer_search_results
from brand_search_client import UserPrompt as BrandUserPrompt, handle_user_prompt_for_Brand, format_brand_search_results
from business_unit_search_client import UserPrompt as BusinessUnitUserPrompt, handle_user_prompt_for_BusinessUnit, format_business_unit_search_results
from coa_search_client import UserPrompt as COAUserPrompt, handle_user_prompt_for_COA, format_coa_search_results
from market_search_client import UserPrompt as MarketUserPrompt, handle_user_prompt_for_Market, format_market_search_results

def perform_subsidiary_search(subsidiary_user_prompt):
    return handle_user_prompt_for_Subsidiary(subsidiary_user_prompt, field="lookup_value")

def perform_customer_search(customer_user_prompt):
    return handle_user_prompt_for_Customer(customer_user_prompt, fields=["lookup_value"])

def perform_brand_search(brand_user_prompt):
    return handle_user_prompt_for_Brand(brand_user_prompt, field="lookup_value")

def perform_business_unit_search(business_unit_user_prompt):
    return handle_user_prompt_for_BusinessUnit(business_unit_user_prompt, field="lookup_value")

def perform_coa_search(coa_user_prompt):
    return handle_user_prompt_for_COA(coa_user_prompt, field="lookup_value")

def perform_market_search(market_user_prompt):
    return handle_user_prompt_for_Market(market_user_prompt, field="lookup_value")

def retrieval_main(user_prompt):
    # Parse user prompt
    with open('config.json', 'r') as f:
        config = json.load(f)

    # Access configuration values
    api_key = config['OPENAI_API_KEY']
    user_info = get_user_info(api_key, user_prompt)

    # CloudSearch integration
    subsidiary_user_prompt = SubsidiaryUserPrompt(query=user_info.subsidiary)
    customer_user_prompt = CustomerUserPrompt(query=user_info.customer)
    brand_user_prompt = BrandUserPrompt(query=user_info.brand)
    business_unit_user_prompt = BusinessUnitUserPrompt(query=user_info.business_unit)
    coa_user_prompt = COAUserPrompt(query=user_info.coa)
    market_user_prompt = MarketUserPrompt(query=user_info.market)

    # Using ThreadPoolExecutor for concurrent execution
    with ThreadPoolExecutor() as executor:
        future_subsidiary = executor.submit(perform_subsidiary_search, subsidiary_user_prompt)
        future_customer = executor.submit(perform_customer_search, customer_user_prompt)
        future_brand = executor.submit(perform_brand_search, brand_user_prompt)
        future_business_unit = executor.submit(perform_business_unit_search, business_unit_user_prompt)
        future_coa = executor.submit(perform_coa_search, coa_user_prompt)
        future_market = executor.submit(perform_market_search, market_user_prompt)

        # Retrieve results
        subsidiary_search_results = future_subsidiary.result()
        customer_search_results = future_customer.result()
        brand_search_results = future_brand.result()
        business_unit_search_results = future_business_unit.result()
        coa_search_results = future_coa.result()
        market_search_results = future_market.result()

    entities_extracted = {
        "customer": user_info.customer,
        "subsidiary": user_info.subsidiary,
        "brand": user_info.brand,
        "business_unit": user_info.business_unit,
        "coa": user_info.coa,
        "market": user_info.market
    }

    return entities_extracted, subsidiary_search_results, customer_search_results, brand_search_results, business_unit_search_results, coa_search_results, market_search_results

if __name__ == "__main__":
    message = "Can you give me the list of accounts for customer Subway, subsidiary Doordash Kitchens and brand Subway Withholding?"
    entities_extracted, subsidiary_search_results, customer_search_results, brand_search_results, business_unit_search_results, coa_search_results, market_search_results = retrieval_main(message)
    print(entities_extracted, subsidiary_search_results, customer_search_results, brand_search_results, business_unit_search_results, coa_search_results, market_search_results)