app.py

from time import sleep
import streamlit as st
import openai
import pinecone
from postgres_db import query_postgresql_realvest

PINECONE_API_KEY = st.secrets["PINECONE_API_KEY"]
OPENAI_API_KEY = st.secrets["OPENAI_API_KEY"]
INDEX_NAME = 'realvest-data-v2'
EMBEDDING_MODEL = "text-embedding-ada-002"  # OpenAI's best embeddings as of Apr 2023
MAX_LENGTH_DESC = 200
MATCH_SCORE_THR = 0.0
TOP_K = 20

def test_pinecone(sleep_time: int=1):
    MAX_TRIALS = 5
    trial = 0
    stats = None
    while (stats is None) and (trial < MAX_TRIALS):
        try:
            print(f"BEFORE: trial: {trial}; stats: {stats}")
            stats = index.describe_index_stats()
            print(f"AFTER: trial: {trial}; stats: {stats}")
            return stats
        except pinecone.core.exceptions.PineconeProtocolError as err:
            print(f"Error, sleep! {err}")
            sleep(sleep_time)
            trial = trial + 1
    
    raise Exception(f'max trials {MAX_TRIALS} Exceeded!')

def query_pinecone(xq, top_k: int=3, include_metadata: bool=True, sleep_time: int=1):
    MAX_TRIALS = 5
    trial = 0
    out = None
    while (out is None) and (trial < MAX_TRIALS):
        try:
            # print(f"BEFORE: trial: {trial}; stats: {out}")
            out = index.query(xq, top_k=top_k, include_metadata=include_metadata)
            # print(f"AFTER: trial: {trial}; stats: {out}")
            return out
        except pinecone.core.exceptions.PineconeProtocolError as err:
            print(f"Error, sleep! {err}")
            sleep(sleep_time)
            trial = trial + 1
    
    raise Exception(f'max trials {MAX_TRIALS} Exceeded!')

def sort_dict_by_value(d: dict, ascending: bool=True):
    """
    Sort dictionary {k1: v1, k2: v2} by its value. The output is
    a sorted list of tuples [(k1, v1), (k2, v2)]
    """
    return sorted(d.items(), key=lambda x: x[1], reverse=not ascending)

def init_session_state():
    if 'display_results' not in st.session_state:
        st.session_state['display_results'] = False

    if 'count_checked' not in st.session_state:
        st.session_state['count_checked'] = 0

def callback_count_checked():
    st.session_state['count_checked'] = 0
    st.session_state['checked_boxes'] = []
    for key in list(st.session_state.keys()):
        if (key.split('__')[0] == 'cb_compare') and (st.session_state[key] == True):
            st.session_state['count_checked'] += 1
            st.session_state['checked_boxes'].append(key)

def summarize_products(products: list) -> str:
    """
    Input:
    products = [
        {text information of product#1},
        {text information of product#2},
        {text information of product#3},
    ]

    Output:
    summary = "{summary of all products}"
    """
    NEW_LINE = '\n'
    prompt = f"""
    Based on the product information below, please read and try to understand it.
    { f"{NEW_LINE*2}---{NEW_LINE*2}".join(products) }
    Please write a concise and insightful summary table (display as HTML) to compare the products for investors, which should inlcude but not limited to:
    - description
    - category
    - asking price
    - location
    - potential profit margin
    """
    print(f"prompt: {prompt}")
    
    openai.api_key = OPENAI_API_KEY
    completion = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ]
    )

    summary = completion.choices[0].message
    return summary

# initialize connection to pinecone (get API key at app.pinecone.io)
pinecone.init(
    api_key=PINECONE_API_KEY,
    environment="us-central1-gcp"  # may be different, check at app.pinecone.io
)
index = pinecone.Index(INDEX_NAME)
# stats = test_pinecone()
# print(f"Pinecone DB stats: {stats}")

### Main
# st.set_page_config(layout="centered")
css='''
<style>
    section.main > div {max-width:70rem}
</style>
'''
st.markdown(css, unsafe_allow_html=True)

# initialize state
init_session_state()

# Create a text input field
query = st.text_input("What are you looking for?")

# Create a button
if st.button('Submit'):

    # initialize
    st.session_state.clear()
    init_session_state()

    # ### call OpenAI text-embedding
    res = openai.Embedding.create(model=EMBEDDING_MODEL, input=[query], api_key=OPENAI_API_KEY)
    xq = res['data'][0]['embedding']
    out = query_pinecone(xq, top_k=TOP_K, include_metadata=True)

    ### candidates
    metadata = {match['metadata']['product_id']: match['metadata'] for match in out['matches']}
    match_score = {match['metadata']['product_id']: match['score'] for match in out['matches']}
    above_thr_sorted = [
        item 
        for item in sort_dict_by_value(match_score, ascending=False) 
        if item[1] > MATCH_SCORE_THR
    ]
    pids = metadata.keys()

    ### query pids
    pids_str = [f"'{pid}'" for pid, _ in above_thr_sorted]
    query = f"""
    SELECT productid, name, category, alternatename, url, logo, description
    FROM main_products
    WHERE productid in ({', '.join(pids_str)});
    """

    results = query_postgresql_realvest(query)
    results = {
        result['productid']: result
        for result in results
    }

    ### For test
    # print(f"above_thr_sorted: {above_thr_sorted}")
    # print(f"results: {results}")
    # print(f"metadata: {metadata}")
    # # TEST ONLY
    # above_thr_sorted = [('2086773', 0.800059378), ('1951083', 0.797319531), ('1998714', 0.795623)]
    # results = {'1951083': {'productid': '1951083', 'name': '2 for 1 Turn-key Business Opportunity in Lynnwood, Washington - BizBuySell', 'category': 'Other', 'alternatename': None, 'url': 'https://www.bizbuysell.com/Business-Opportunity/2-for-1-Turn-key-Business-Opportunity/1951083/', 'logo': 'https://images.bizbuysell.com/shared/listings/195/1951083/87198e08-a191-4d97-a33b-9e9f40fa02f4-W768.jpg', 'description': 'Your chance to own a successful Korean traditional KBBQ grill restaurant and Korean dive bar. Owner is retiring after 19 years of business. This Korean BBQ restaurant utilizes a traditional grill called "Sot Ttu Kkeong" widely found in Korea. There are 10 separate grilling tables with a unique hood system to eliminate odors immediately. The bar next door may be able to extend hours into the summer. With one shared full kitchen, the new owner will be able to maximize business and potentially earn double income.'}, '1998714': {'productid': '1998714', 'name': 'Portland CPA Firm in Portland, Oregon - BizBuySell', 'category': 'Accounting and Tax Practices', 'alternatename': None, 'url': 'https://www.bizbuysell.com/Business-Opportunity/Portland-CPA-Firm/1998714/', 'logo': 'https://images.bizbuysell.com/shared/listings/199/1998714/cd02bdb9-32c9-409d-b82e-d0531c12eb39-W768.jpg', 'description': 'OR1002: UPDATED :The seller of this Portland CPA firm is approaching retirement and ready to sell the firm. The firm has a great reputation, has good systems in place, is paperless, and has a great staff. The mix of services offers a consistent stream of cash flow to the owner. The seller is seeking a CPA buyer. The office space is available for continued lease after the sale. Revenues for sale include:7% Accounting, bookkeeping and payroll services26% Income tax preparation services for individual clients35% Income tax preparation services for business and other clients28% Audits and reviews4% Consulting services'}, '2086773': {'productid': '2086773', 'name': 'Asian Grocery Supermarket, 1 owner for 29 years in Salem, Oregon - BizBuySell', 'category': 'Grocery Stores and Supermarkets', 'alternatename': None, 'url': 'https://www.bizbuysell.com/Business-Real-Estate-For-Sale/Asian-Grocery-Supermarket-1-owner-for-29-years/2086773/', 'logo': 'https://images.bizbuysell.com/shared/listings/208/2086773/861f6ba6-a994-4e90-9c62-0a593dae2a31-W768.jpg', 'description': 'Great location, well established and profitable supermarket.We have been the sole owner for almost 29 years, so business boasts of a great reputation.'}}
    # metadata = {'2086773': {'asking_price': 1000000.0, 'asking_price_currency': 'USD', 'building_status': 'Established', 'category': 'Grocery Stores and Supermarkets', 'chunk_type': 'profile', 'city': 'Salem', 'document': '# Listing Profile\n    \nAsking Price (USD): 1000000            \n\nReason for Selling: Retire                                                                ', 'listing_type': 'Retail', 'location': 'Salem, OR', 'main_category': 'Grocery Stores and Supermarkets', 'offer_type': 'Offer', 'offers__available_from__address__locality': 'Salem', 'offers__available_from__address__region': 'Oregon', 'offers__available_from__address__type': 'PostalAddress', 'offers__available_from__type': 'Place', 'product_id': '2086773', 'similar_pids': ['2074401', '2087795', '2068650'], 'state_code': 'OR'}, '1951083': {'asking_price': 200000.0, 'asking_price_currency': 'USD', 'category': 'Other', 'chunk_type': 'profile', 'city': 'Lynnwood', 'document': '# Listing Profile\n    \nAsking Price (USD): 200000            \n\nReason for Selling: Retiring                                                                ', 'location': 'Lynnwood, WA', 'main_category': 'Other', 'offer_type': 'Offer', 'offers__available_from__address__locality': 'Lynnwood', 'offers__available_from__address__region': 'Washington', 'offers__available_from__address__type': 'PostalAddress', 'offers__available_from__type': 'Place', 'product_id': '1951083', 'similar_pids': ['2113741', '2033980', '2034855'], 'state_code': 'WA'}, '1998714': {'asking_price': 900000.0, 'asking_price_currency': 'USD', 'category': 'Accounting and Tax Practices', 'chunk_type': 'profile', 'city': 'Portland', 'document': '# Listing Profile\n    \nAsking Price (USD): 900000            \n\nReason for Selling: Approaching retirement                                                                ', 'fin__gross_revenue': 958000.0, 'location': 'Portland, OR', 'main_category': 'Accounting and Tax Practices', 'offer_type': 'Offer', 'offers__available_from__address__locality': 'Portland', 'offers__available_from__address__region': 'Oregon', 'offers__available_from__address__type': 'PostalAddress', 'offers__available_from__type': 'Place', 'product_id': '1998714', 'similar_pids': ['2026155', '2066311'], 'state_code': 'OR'}}    

    # update
    st.session_state['above_thr_sorted'] = above_thr_sorted
    st.session_state['results'] = results
    st.session_state['metadata'] = metadata

    st.session_state['display_results'] = True

if st.session_state['display_results']:

    summary_container = st.empty()

    st.header("Results")
    st.divider()

    # display matched results
    for pid, match_score in st.session_state['above_thr_sorted']:

        if pid not in st.session_state['results']:
            continue

        metadata_pid = st.session_state['metadata'][pid]
        result = st.session_state['results'][pid]
        col_icon, col_info, col_compare = st.columns([2, 6, 1])

        with col_icon:
            st.image(result["logo"])

        with col_info:
            # TODO: make asking price display $xxx,xxx
            st.markdown(f"""match score: { round(100 * match_score, 2) }
            <br>
            **{result['name']}**
            <br>
            _Asking Price:_ {metadata_pid.get('asking_price', 'N/A')}
            <br>
            _Category:_ {metadata_pid.get('category', 'N/A')}
            <br>
            _Location:_ {metadata_pid.get('location', 'N/A')}
            """, unsafe_allow_html=True)

            st.markdown(f"""**_Description:_** {result['description'][:MAX_LENGTH_DESC]}...[more]({result['url']})
            """)

        with col_compare:
            st.checkbox('compare', key=f"cb_compare__{pid}", on_change=callback_count_checked)

# display summary tab
if st.session_state['count_checked'] > 0:
    with summary_container.container():
        st.header('Summary')
        if st.button('Compare Products'):
            products = []
            for key in st.session_state['checked_boxes']:
                # TODO: Need to pull all the document
                # TODO: Need to dedup the pid too
                pid = key.split('__')[-1]
                products.append(
                    st.session_state['metadata'][pid].get('document')
                )
            summary = summarize_products(products)
            st.markdown(summary.get("content"), unsafe_allow_html=True)
else:
    try:
        summary_container.empty()
    except NameError:
        pass


with st.expander("developer tool"):
    st.json(st.session_state)