Real Estate Natural Language Querying Agent

# Text similarity space for addresses
address_space = sl.TextSimilaritySpace(
    text=real_estate.address,
    model="sentence-transformers/all-mpnet-base-v2"
)

# Numerical similarity spaces
price_space = sl.NumberSpace(
    number=real_estate.price,
    min_value=0,
    max_value=int(max_price),
    mode=sl.Mode.MINIMUM
)

beds_space = sl.NumberSpace(
    number=real_estate.number_beds,
    min_value=0,
    max_value=int(max_beds),
    mode=sl.Mode.MINIMUM
)

baths_space = sl.NumberSpace(
    number=real_estate.number_baths,
    min_value=0,
    max_value=int(max_baths),
    mode=sl.Mode.MINIMUM
)

population_space = sl.NumberSpace(
    number=real_estate.population,
    min_value=0,
    max_value=int(max_population),
    mode=sl.Mode.MINIMUM
)

income_space = sl.NumberSpace(
    number=real_estate.median_family_income,
    min_value=0,
    max_value=int(max_median_family_income),
    mode=sl.Mode.MINIMUM
)

# Categorical similarity spaces
city_space = sl.CategoricalSimilaritySpace(
    category_input=real_estate.city,
    categories=unique_city_categories
)

province_space = sl.CategoricalSimilaritySpace(
    category_input=real_estate.province,
    categories=unique_province_categories,
    negative_filter=-1.0,
    uncategorized_as_category=False
)

query = sl.Query(index, weights={
    address_space: sl.Param("address_weight", default=1.0),
    price_space: sl.Param("price_weight", default=100.0),
    city_space: sl.Param("city_weight", default=1.0),
    beds_space: sl.Param("beds_weight", default=1.0),
    baths_space: sl.Param("baths_weight", default=1.0),
    population_space: sl.Param("population_weight", default=0.2),
    income_space: sl.Param("income_weight", default=1.0),
})

class RealEstateAgent:
    def __init__(self, df, app, master_query, openai_client):
        self.tools = [
            PropertyRetrievalTool(),      # Basic "show me homes" requests
            PropertyRecommendationTool(), # "What's good for me?" queries
            QueryRefinementTool(),        # When searches come up empty
            NarrativeInsightTool(),       # Market analysis stuff
            MultiStepQueryTool()          # "Compare Toronto vs Vancouver"
        ]

class PropertyRetrievalTool(Tool):
    def name(self) -> str:
        return "PropertyRetrievalTool"

    def description(self) -> str:
        return "Retrieves properties based on natural language queries, leveraging NLQ intelligence for ranking"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        try:
            results = app.query(master_query, natural_query=query, limit=10)
            df_results = results_to_dataframe(results, df)
            if len(df_results) == 0:
                # If there are no results found, pass it to the QueryRefinementTool to get the better response
                refine_tool = QueryRefinementTool()
                return refine_tool.use(query, df, app, master_query, openai_client)
            return f"Search Results for '{query}':\\n\\n{format_property_display(df_results, query, detailed=True)}"
        except Exception as e:
            return f"Error in property retrieval: {str(e)}"

class PropertyRecommendationTool(Tool):
    def name(self) -> str:
        return "PropertyRecommendationTool"

    def description(self) -> str:
        return "Provides personalized property recommendations based on inferred user preferences"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        # Infer user profile from query
        profile_prompt = f"""
        Extract user preferences from this query, such as budget, family size, commute preferences, or lifestyle (e.g., urban, family-friendly).
        Query: "{query}"
        Return a JSON object with inferred preferences (e.g., budget). If none, return .
        """
        try:
            profile_response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": profile_prompt}],
                temperature=0.3,
                max_tokens=100
            )
            user_profile = eval(profile_response.choices[0].message.content.strip()) or {}
        except Exception as e:
            user_profile = {}
            print(f"Error inferring user profile: {str(e)}")

        # Fetch properties using NLQ
        results = app.query(master_query, natural_query=query, limit=15)
        df_results = results_to_dataframe(results, df)
        if len(df_results) == 0:
            return "No properties found to recommend."

        # Prepare property data (top properties from NLQ ranking)
        properties_summary = [
            {
                'address': row['address'],
                'city': row['city'],
                'province': row['province'],
                'price': int(row['price']),
                'beds': int(row['number_beds']),
                'baths': int(row['number_baths']),
                'population': int(row['population']),
                'median_income': int(row['median_family_income']),
                'relevance': round(row['relevance_score'], 3)
            }
            for _, row in df_results.head(10).iterrows()
        ]

        # Generate recommendations
        prompt = f"""
        You are a real estate advisor. Provide recommendations based on:
        Query: "{query}"
        Inferred Profile: {user_profile}
        Properties (ranked by NLQ relevance): {properties_summary}

        The properties are already ranked by our intelligent NLQ system considering all factors including price relevance.
        Provide:
        1. Top 3 recommendations with reasons (include address) - respect the NLQ ranking
        2. Pros and cons for each
        3. Additional considerations (e.g., neighborhood, lifestyle fit)
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=800
            )
            recommendations = response.choices[0].message.content.strip()

            # Use top properties from NLQ ranking (no additional sorting)
            formatted_properties = format_property_display(df_results.head(5), query, detailed=True)
            return f"Recommendations for '{query}':\\n\\n{recommendations}\\n\\nProperty Listings (NLQ Ranked):\\n{formatted_properties}"
        except Exception as e:
            return f"Error generating recommendations: {str(e)}"

class QueryRefinementTool(Tool):
    def name(self) -> str:
        return "QueryRefinementTool"

    def description(self) -> str:
        return "Refines queries when results are sparse or irrelevant"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        # Try initial query
        results = app.query(master_query, natural_query=query, limit=10)
        df_results = results_to_dataframe(results, df)

        if len(df_results) >= 3:
            return f"Search Results for '{query}':\\n\\n{format_property_display(df_results, query, detailed=True)}"

        # Refine query if fewer than 3 results
        refine_prompt = f"""
        The query "{query}" returned {len(df_results)} results. Suggest an alternative query to get more relevant results.
        Consider relaxing filters (e.g., price, location) or expanding scope (e.g., nearby cities).
        Return a single alternative query as a string.
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": refine_prompt}],
                temperature=0.3,
                max_tokens=50
            )
            new_query = response.choices[0].message.content.strip()

            # Try refined query
            results = app.query(master_query, natural_query=new_query, limit=10)
            df_new_results = results_to_dataframe(results, df)

            if len(df_new_results) == 0:
                return f"No properties found for '{query}'. Suggested query '{new_query}' also returned no results."

            return f"""
Original Query: '{query}' returned {len(df_results)} results.
Suggested Query: '{new_query}'
Results for '{new_query}' (NLQ Ranked):
{format_property_display(df_new_results, new_query, detailed=True)}
            """
        except Exception as e:
            return f"Error refining query: {str(e)}"

class NarrativeInsightTool(Tool):
    def name(self) -> str:
        return "NarrativeInsightTool"

    def description(self) -> str:
        return "Provides narrative insights for complex queries (e.g., investment potential)"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        results = app.query(master_query, natural_query=query, limit=10)
        df_results = results_to_dataframe(results, df)
        if len(df_results) == 0:
            return "No properties found for analysis."

        stats = {
            'avg_price': int(df_results['price'].mean()),
            'median_price': int(df_results['price'].median()),
            'price_range': (int(df_results['price'].min()), int(df_results['price'].max())),
            'avg_beds': round(df_results['number_beds'].mean(), 1),
            'avg_baths': round(df_results['number_baths'].mean(), 1),
            'cities': df_results['city'].value_counts().head(3).to_dict(),
            'avg_income': int(df_results['median_family_income'].mean())
        }

        prompt = f"""
        Provide a narrative analysis for the query: "{query}"
        Data: {stats}
        Properties (NLQ ranked by relevance): {[{k: v for k, v in row.items() if k in ['address', 'city', 'price', 'number_beds', 'number_baths', 'relevance_score']} for _, row in df_results.head(5).iterrows()]}

        Note: Properties are already intelligently ranked by our NLQ system considering price and other factors.
        Include:
        1. Market overview
        2. Investment potential (if relevant)
        3. Key considerations (e.g., location, price trends)
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=600
            )
            insights = response.choices[0].message.content.strip()
            formatted_properties = format_property_display(df_results, query, detailed=True)
            return f"Insights for '{query}':\\n\\n{insights}\\n\\nProperty Listings (NLQ Ranked):\\n{formatted_properties}"
        except Exception as e:
            return f"Error generating insights: {str(e)}"

class MultiStepQueryTool(Tool):
    def name(self) -> str:
        return "MultiStepQueryTool"

    def description(self) -> str:
        return "Handles queries requiring multiple NLQ calls (e.g., comparisons)"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        # Detect if query involves comparison
        compare_pattern = re.compile(r'\\b(compare|versus|vs\\.?)\\b', re.IGNORECASE)
        if not compare_pattern.search(query):
            return "Query does not require comparison."

        # Extract sub-queries
        prompt = f"""
        Split the query into sub-queries for comparison.
        Query: "{query}"
        Return a JSON list of sub-queries (e.g., ["3 bedroom homes in Toronto", "3 bedroom homes in Kitchener"]).
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=100
            )
            sub_queries = eval(response.choices[0].message.content.strip())
            if not isinstance(sub_queries, list) or len(sub_queries) < 2:
                return "Unable to split query for comparison."

            # Run sub-queries (NLQ handles intelligent ranking)
            results_dict = {}
            for sub_query in sub_queries[:2]:  # Limit to 2 for simplicity
                results = app.query(master_query, natural_query=sub_query, limit=5)
                df_sub_results = results_to_dataframe(results, df)
                results_dict[sub_query] = df_sub_results

            # Generate comparison
            data_str = ""
            for sub_query, df_sub_results in results_dict.items():
                data_str += f"\\n{sub_query} (NLQ ranked):\\n"
                for _, row in df_sub_results.head(3).iterrows():
                    data_str += f"  - Address: {row['address']}, City: {row['city']}, Price: ${row['price']:,.0f}, Beds: {row['number_beds']}, Baths: {row['number_baths']}, Relevance: {row['relevance_score']:.3f}\\n"

            compare_prompt = f"""
            Compare properties from these sub-queries (each ranked by NLQ intelligence):
            {list(results_dict.keys())}
            Data:
            {data_str}

            Note: Each property list is already ranked by our NLQ system considering all factors including price relevance.
            Provide:
            1. Side-by-side comparison (price, beds, baths, NLQ relevance)
            2. Key differences (e.g., location, value, market positioning)
            3. Recommendations for different buyer types
            """
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": compare_prompt}],
                temperature=0.7,
                max_tokens=800
            )
            comparison = response.choices[0].message.content.strip()

            # Format results
            output = [f"Comparison for '{query}':\\n\\n{comparison}"]
            for sub_query, df_sub_results in results_dict.items():
                if not df_sub_results.empty:
                    output.append(f"\\nResults for '{sub_query}' (NLQ Ranked):\\n{format_property_display(df_sub_results, sub_query, detailed=True)}")

            return "\\n".join(output)
        except Exception as e:
            return f"Error processing comparison: {str(e)}"

def classify_query(self, query: str) -> str:
        valid_categories = {'retrieval', 'recommendation', 'refinement', 'insight', 'multistep'}
        prompt = f"""
Classify the query into one category based on its intent, returning only the category name ('retrieval' if unsure):
- retrieval: Search for properties by criteria (e.g., location, price, bedrooms).
- recommendation: Personalized property suggestions for user needs.
- refinement: Strict criteria likely yielding few results, needing adjustment.
- insight: Analysis, investment advice, or market trends.
- multistep: Comparison or multiple criteria across locations.
Query: "{query}"
        """
        try:
            response = self.openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=20
            )
            classification = response.choices[0].message.content.strip().lower()
            return classification if classification in valid_categories else 'retrieval'
        except Exception as e:
            print(f"Classification error: {str(e)}")
            return 'retrieval'