screens/search.py

import streamlit as st
import os
import streamlit.components.v1 as components
from io import BytesIO
import requests
import ast

from langchain import PromptTemplate
from langchain.chains import RetrievalQA
from langchain.vectorstores import Chroma
from langchain.embeddings import SentenceTransformerEmbeddings
from bardapi import Bard
from typing import Any, List, Mapping, Optional

import yaml
with open("config.yml", "r") as ymlfile:
    cfg = yaml.safe_load(ymlfile)
os.environ['_BARD_API_KEY'] = cfg["API_KEY"]["Bard"]

from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun
class BardLLM(LLM):


    @property
    def _llm_type(self) -> str:
        return "custom"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
    ) -> str:
        response = Bard(token=os.environ['_BARD_API_KEY']).get_answer(prompt)['content']
        return response

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {}
    
@st.cache_data
def get_image(url):
    r = requests.get(url)
    return BytesIO(r.content)


# Define global variables
embeddings = None
index = None
QUESTION_PROMPT = None
qa = None
result = []

# Custom session state class for managing pagination
class SessionState:
    def __init__(self):
        self.page_index = 0  # Initialize page index
        self.database_loaded = False  # Initialize database loaded state

# Create a session state object
session_state = SessionState()

# Define the search function outside of Search_Property
def display_search_results(result, start_idx, end_idx):
    if result:
        st.subheader("Search Results:")
        for idx in range(start_idx, end_idx):
            if idx >= len(result):
                break
            property_info = result[idx]
            st.markdown(f"**Result {idx + 1}**")
            
            # Display property information
            if 'Image URL' in property_info.metadata and property_info.metadata['Image URL'] is not None and not isinstance(property_info.metadata['Image URL'], float):
                image_path_urls = property_info.metadata['Image URL']
                if image_path_urls is not None and not isinstance(image_path_urls, float):
                    # Convert the string to a Python list
                    imageUrls = ast.literal_eval(image_path_urls)

                    # Now, imageUrls is a list of strings
                    st.image(imageUrls[0],width=700)

            st.markdown(f"🏡 {property_info.metadata['Title']}")
            if 'Location' in property_info.metadata and property_info.metadata['Location'] is not None and not isinstance(property_info.metadata['Location'], float):
                st.write(f"📍 Address: {property_info.metadata['Location']}")
            if 'Area' in property_info.metadata and property_info.metadata['Area'] is not None and not isinstance(property_info.metadata['Area'], float):
                st.markdown(f"📏 Size: {property_info.metadata['Area']}")
            if 'Price' in property_info.metadata and property_info.metadata['Price'] is not None and not isinstance(property_info.metadata['Price'], float):
                st.markdown(f"💰 Price: {property_info.metadata['Price']} ")
            st.markdown(f"📅 Published Date: {property_info.metadata['Time stamp']}") 
            col3, col4 = st.columns([2, 1]) 
            with col3: 
                with st.expander("Full Property Information"):
                    st.write(f"🏡 Property Title: {property_info.metadata['Title']}")
                    if 'Area' in property_info.metadata and property_info.metadata['Area'] is not None and not isinstance(property_info.metadata['Area'], float):
                        st.write(f"📏 Size: {property_info.metadata['Area']}")
                    if 'Category' in property_info.metadata and property_info.metadata['Category'] is not None and not isinstance(property_info.metadata['Category'], float):
                        st.write(f"🏢 Category: {property_info.metadata['Category']}")
                    if 'Description' in property_info.metadata and property_info.metadata['Description'] is not None and not isinstance(property_info.metadata['Description'], float):
                        st.write(f"📝 Description: {property_info.metadata['Description']}")
                    if 'Price' in property_info.metadata and property_info.metadata['Price'] is not None and not isinstance(property_info.metadata['Price'], float):
                        st.write(f"💰 Price: {property_info.metadata['Price']}")
                    st.write(f"📅 Date: {property_info.metadata['Time stamp']}")
                    if 'Location' in property_info.metadata and property_info.metadata['Location'] is not None and not isinstance(property_info.metadata['Location'], float):
                        st.write(f"📍 Address: {property_info.metadata['Location']}")
                    st.write(f"🆔 ID: {property_info.metadata['ID']}")
                    if 'Estate type' in property_info.metadata and property_info.metadata['Estate type'] is not None and not isinstance(property_info.metadata['Estate type'], float):
                        st.write(f"🏠 Housing Type: {property_info.metadata['Estate type']}")
                    if 'Email' in property_info.metadata and property_info.metadata['Email'] is not None and not isinstance(property_info.metadata['Email'], float):
                        st.write(f"✉️ Email: {property_info.metadata['Email']}")
                    if 'Mobile Phone' in property_info.metadata and property_info.metadata['Mobile Phone'] is not None and not isinstance(property_info.metadata['Mobile Phone'], float):
                        st.write(f"📞 Phone: {property_info.metadata['Mobile Phone']}")
                    if 'Certification status' in property_info.metadata and property_info.metadata['Certification status'] is not None and not isinstance(property_info.metadata['Certification status'], float):
                        st.write(f"🏆 Certification status: {property_info.metadata['Certification status']}")
                    if 'Direction' in property_info.metadata and property_info.metadata['Direction'] is not None and not isinstance(property_info.metadata['Direction'], float):
                        st.write(f"🧭 Direction: {property_info.metadata['Direction']}")
                    if 'Rooms' in property_info.metadata and property_info.metadata['Rooms'] is not None and not isinstance(property_info.metadata['Rooms'], float):
                        st.write(f"🚪 Rooms: {property_info.metadata['Rooms']}")
                    if 'Bedrooms' in property_info.metadata and property_info.metadata['Bedrooms'] is not None and not isinstance(property_info.metadata['Bedrooms'], float):
                        st.write(f"🛏️ Bedrooms: {property_info.metadata['Bedrooms']}")
                    if 'Kitchen' in property_info.metadata and property_info.metadata['Kitchen'] is not None and not isinstance(property_info.metadata['Kitchen'], float):
                        st.write(f"🍽️ Kitchen: {property_info.metadata['Kitchen']}")
                    if 'Living room' in property_info.metadata and property_info.metadata['Living room'] is not None and not isinstance(property_info.metadata['Living room'], float):
                        st.write(f"🛋️ Living room: {property_info.metadata['Living room']}")
                    if 'Bathrooms' in property_info.metadata and property_info.metadata['Bathrooms'] is not None and not isinstance(property_info.metadata['Bathrooms'], float):
                        st.write(f"🚽 Bathrooms: {property_info.metadata['Bathrooms']}")
                    if 'Front width' in property_info.metadata and property_info.metadata['Front width'] is not None and not isinstance(property_info.metadata['Front width'], float):
                        st.write(f"📐 Front width: {property_info.metadata['Front width']}")
                    if 'Floor' in property_info.metadata and property_info.metadata['Floor'] is not None and not isinstance(property_info.metadata['Floor'], float):
                        st.write(f"🧱 Floor: {property_info.metadata['Floor']}")
                    if 'Parking Slot' in property_info.metadata and property_info.metadata['Parking Slot'] is not None and not isinstance(property_info.metadata['Parking Slot'], float):
                        st.write(f"🚗 Parking Slot: {property_info.metadata['Parking Slot']}")
                    if 'Seller name' in property_info.metadata and property_info.metadata['Seller name'] is not None and not isinstance(property_info.metadata['Seller name'], float):
                        st.write(f"👤 Seller Name: {property_info.metadata['Seller name']}")
                    if 'Seller type' in property_info.metadata and property_info.metadata['Seller type'] is not None and not isinstance(property_info.metadata['Seller type'], float):
                        st.write(f"👨‍💼 Seller type: {property_info.metadata['Seller type']}")
                    if 'Seller Address' in property_info.metadata and property_info.metadata['Seller Address'] is not None and not isinstance(property_info.metadata['Seller Address'], float):
                        st.write(f"📌 Seller Address: {property_info.metadata['Seller Address']}")
                    if 'Balcony Direction' in property_info.metadata and property_info.metadata['Balcony Direction'] is not None and not isinstance(property_info.metadata['Balcony Direction'], float):
                        st.write(f"🌄 Balcony Direction: {property_info.metadata['Balcony Direction']}")
                    if 'Furniture' in property_info.metadata and property_info.metadata['Furniture'] is not None and not isinstance(property_info.metadata['Furniture'], float):
                        st.write(f"🛋️ Furniture: {property_info.metadata['Furniture']}")
                    if 'Toilet' in property_info.metadata and property_info.metadata['Toilet'] is not None and not isinstance(property_info.metadata['Toilet'], float):
                        st.write(f"🚽 Toilet: {property_info.metadata['Toilet']}")                

            with col4:
                st.empty()
            if 'Image URL' in property_info.metadata and property_info.metadata['Image URL'] is not None and not isinstance(property_info.metadata['Image URL'], float):
                imageCarouselComponent = components.declare_component("image-carousel-component", path="./frontend/public")
                image_path_urls = property_info.metadata['Image URL']
                if image_path_urls is not None and not isinstance(image_path_urls, float):
                    # Convert the string to a Python list
                    imageUrls = ast.literal_eval(image_path_urls)
                    if len(imageUrls) > 1:
                        selectedImageUrl = imageCarouselComponent(imageUrls=imageUrls, height=200)
                        if selectedImageUrl is not None:
                            st.image(selectedImageUrl)

            # Add a divider after displaying property info
            st.markdown("<hr style='border: 2px solid white'>", unsafe_allow_html=True)  # Horizontal rule as a divider
            

def Search_Property():
    global embeddings, index, result, QUESTION_PROMPT, qa

    st.title("🏘️ Property Search ")
    # Load data and create the search
    if not session_state.database_loaded:
        st.info("Loading database... This may take a moment.")
        embeddings = SentenceTransformerEmbeddings(model_name="keepitreal/vietnamese-sbert")
        # Create a Chroma object with persistence
        db = Chroma(persist_directory="./chroma_index_1", embedding_function=embeddings)
        # Get documents from the database
        db.get()
        llm=BardLLM()
        qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=db.as_retriever(search_type="similarity", search_kwargs={"k":4}),
        return_source_documents=True)
        question_template = """
        Context: You are a helpful and informative bot that answers questions posed below using provided context.\
        You have to be truthful. Do not recommend or propose any infomation of the properties.\
        Be sure to respond in a complete sentence, being comprehensive, including all information in the provided context.\
        Imagine you're talking to a friend and use natural language and phrasing.\
        You can only use Vietnamese do not use other languages.

        QUESTION: '{question}'

        ANSWER:
        """
        QUESTION_PROMPT = PromptTemplate(
            template=question_template, input_variables=["question"]
        )   
        session_state.database_loaded = True

    if session_state.database_loaded:
        col1, col2 = st.columns([2, 1])  # Create a two-column layout

        with col1:
            query = st.text_input("Enter your property search query:")
            search_button = st.button("Search", help="Click to start the search")

            if search_button:
                if not query:
                    st.warning("Please input your query")
                else:
                    with st.spinner("Searching..."):
                        if query is not None:  # Check if model_embedding is not None
                            qa.combine_documents_chain.llm_chain.prompt = QUESTION_PROMPT
                            qa.combine_documents_chain.verbose = True
                            qa.return_source_documents = True
                            results = qa({"query":query,})
                            result = results["source_documents"]
                            session_state.page_index = 0  # Reset page index when a new search is performed
                        
        with col2:
            if len(result) > 0:
                st.info(f'Total Results: {len(result)} properties found.')  # Display "Total Results" in the second column

        if result:   
            N = 5 
            prev_button, next_button = st.columns([4,1])
            last_page = len(result) // N

            
            # Update page index based on button clicks
            if prev_button.button("Previous", key="prev_button"):
                if session_state.page_index - 1 < 0:
                    session_state.page_index = last_page
                else:
                    session_state.page_index -= 1

            if next_button.button("Next", key="next_button"):
                if session_state.page_index > last_page:
                    st.warning("Displayed all results")
                    session_state.page_index = 0
                else:
                    session_state.page_index += 1

            # Calculate the range of results to display (5 properties at a time)
            start_idx = session_state.page_index * N 
            end_idx = (1 + session_state.page_index) * N

            # Display results for the current page
            display_search_results(result, start_idx, end_idx)