pathakDev10/EstateGuru

1import uuid2import threading3import asyncio4import json5import re6from datetime import datetime7from fastapi import FastAPI, WebSocket, WebSocketDisconnect8from langchain_core.messages import AIMessage, HumanMessage, SystemMessage9from langgraph.graph import StateGraph, START, END10import faiss11from sentence_transformers import SentenceTransformer12import pickle13import numpy as np14from tools import extract_json_from_response, apply_filters_partial, rule_based_extract, format_property_data, estateKeywords15import random16from langchain_core.prompts import ChatPromptTemplate17from langchain_core.tools import tool18from langchain_core.callbacks import StreamingStdOutCallbackHandler, CallbackManager19from langchain_core.callbacks.base import BaseCallbackHandler20from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer21 22 23class CallbackTextStreamer(TextStreamer):24    def __init__(self, tokenizer, callbacks, skip_prompt=True, skip_special_tokens=True):25        super().__init__(tokenizer, skip_prompt=skip_prompt, skip_special_tokens=skip_special_tokens)26        self.callbacks = callbacks27 28    def on_new_token(self, token: str):29        for callback in self.callbacks:30            callback.on_llm_new_token(token)31 32 33 34 35 36class ChatQwen:37    def __init__(self, temperature=0.3, streaming=False, max_new_tokens=512, callbacks=None):38        self.temperature = temperature39        self.streaming = streaming40        self.max_new_tokens = max_new_tokens41        self.callbacks = callbacks42        self.model_name = "Qwen/Qwen2.5-1.5B-Instruct"43        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)44        self.model = AutoModelForCausalLM.from_pretrained(45            self.model_name, 46            torch_dtype="auto", 47            device_map="auto"48        )49 50    def generate_text(self, messages: list) -> str:51        """52        Given a list of messages, create a prompt and generate text using the Qwen model.53        In streaming mode, uses a TextIteratorStreamer and iterates over tokens to call callbacks.54        """55        # Create prompt from messages using the tokenizer's chat template.56        prompt = self.tokenizer.apply_chat_template(57            messages,58            tokenize=False,59            add_generation_prompt=True60        )61        model_inputs = self.tokenizer([prompt], return_tensors="pt").to(self.model.device)62 63        if self.streaming:64            from transformers import TextIteratorStreamer65            from threading import Thread66 67            # Create the streamer that collects tokens as they are generated.68            streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)69            generation_kwargs = dict(70                **model_inputs,71                max_new_tokens=self.max_new_tokens,72                streamer=streamer,73                temperature=self.temperature,74                do_sample=True75            )76            # Run generation in a separate thread so that we can iterate over tokens.77            thread = Thread(target=self.model.generate, kwargs=generation_kwargs)78            thread.start()79 80            generated_text = ""81            # Iterate over tokens as they arrive.82            for token in streamer:83                generated_text += token84                # Call each callback with the new token.85                if self.callbacks:86                    for callback in self.callbacks:87                        callback.on_llm_new_token(token)88            # In streaming mode you may want to return empty string,89            # but here we return the full text if needed.90            return generated_text91        else:92            outputs = self.model.generate(93                **model_inputs,94                max_new_tokens=self.max_new_tokens,95                temperature=self.temperature,96                do_sample=True97            )98            # Remove the prompt tokens from the output.99            prompt_length = model_inputs.input_ids.shape[-1]100            generated_ids = outputs[0][prompt_length:]101            text_output = self.tokenizer.decode(generated_ids, skip_special_tokens=True)102            return text_output103 104 105    def invoke(self, messages: list, config: dict = None) -> AIMessage:106        config = config or {}107        # Use provided callbacks if any, otherwise default to the callbacks in the instance.108        callbacks = config.get("callbacks", self.callbacks)109        original_callbacks = self.callbacks110        self.callbacks = callbacks111 112        output_text = self.generate_text(messages)113        self.callbacks = original_callbacks114 115        if self.streaming:116            return AIMessage(content="")117        else:118            return AIMessage(content=output_text)119 120 121    def __call__(self, messages: list) -> AIMessage:122        return self.invoke(messages)123 124 125 126class WebSocketStreamingCallbackHandler(BaseCallbackHandler):127    def __init__(self, connection_id: str, loop):128        self.connection_id = connection_id129        self.loop = loop130 131    def on_llm_new_token(self, token: str, **kwargs):132        asyncio.run_coroutine_threadsafe(133            manager_socket.send_message(self.connection_id, token),134            self.loop135        )136 137 138llm = ChatQwen(temperature=0.3, streaming=True, max_new_tokens=512)139 140index = faiss.read_index("./faiss.index")141with open("./metadata.pkl", "rb") as f:142    docs = pickle.load(f)143st_model = SentenceTransformer('all-MiniLM-L6-v2')144 145 146def make_system_prompt(suffix: str) -> str:147    return (148        "You are EstateGuru, a real estate expert created by Abhishek Pathak from SwavishTek. "149        "Your role is to help customers buy properties using the available data. "150        "Only use the provided data—do not make up any information. "151        "The default currency is AED. If a query uses a different currency, convert the amount to AED "152        "(for example, $10k becomes 36726.50 AED and $1 becomes 3.67 AED). "153        "If a customer is interested in a property, wants to buy, or needs to contact an agent or customer care, "154        "instruct them to call +91 8766268285."155        f"\n{suffix}"156    )157 158general_query_prompt = make_system_prompt(159    "You are EstateGuru, a helpful real estate assistant. Answer the user's query accurately using the available data. "160    "Do not invent any details or go beyond the real estate domain. "161    "If the user shows interest in a property or contacting an agent, ask them to call +91 8766268285."162)163 164# ------------------------ Tool Definitions ------------------------165 166@tool167def extract_filters(query: str) -> dict:168    """For extracting filters"""169    # Use a non-streaming ChatQwen for tool use.170    llm_local = ChatQwen(temperature=0.3, streaming=False)171    system = (172        "You are an expert in extracting filters from property-related queries. Your task is to extract and return only the keys explicitly mentioned in the query as a valid JSON object (starting with '{' and ending with '}'). Include only those keys that are directly present in the query.\n\n"173        "The possible keys are:\n"174        "  - 'projectName': The name of the project.\n"175        "  - 'developerName': The developer's name.\n"176        "  - 'relationshipManager': The relationship manager.\n"177        "  - 'propertyAddress': The property address.\n"178        "  - 'surroundingArea': The area or nearby landmarks.\n"179        "  - 'propertyType': The type or configuration of the property.\n"180        "  - 'amenities': Any amenities mentioned.\n"181        "  - 'coveredParking': Parking availability.\n"182        "  - 'petRules': Pet policies.\n"183        "  - 'security': Security details.\n"184        "  - 'occupancyRate': Occupancy information.\n"185        "  - 'constructionImpact': Construction or its impact.\n"186        "  - 'propertySize': Size of the property.\n"187        "  - 'propertyView': View details.\n"188        "  - 'propertyCondition': Condition of the property.\n"189        "  - 'serviceCharges': Service or maintenance charges.\n"190        "  - 'ownershipType': Ownership type.\n"191        "  - 'totalCosts': A cost threshold or cost amount.\n"192        "  - 'paymentPlans': Payment or financing plans.\n"193        "  - 'expectedRentalYield': Expected rental yield.\n"194        "  - 'rentalHistory': Rental history.\n"195        "  - 'shortTermRentals': Short-term rental information.\n"196        "  - 'resalePotential': Resale potential.\n"197        "  - 'uniqueId': A unique identifier.\n\n"198        "Important instructions regarding cost thresholds:\n"199        "  - If the query contains phrases like 'under 10k', 'below 2m', or 'less than 5k', interpret these as cost thresholds.\n"200        "  - Convert any shorthand cost values to pure numbers (for example, '10k' becomes 10000, '2m' becomes 2000000) and assign them to the key 'totalCosts'.\n"201        "  - Do not use 'propertySize' for cost thresholds.\n\n"202        "  - Default currency is AED, if user query have different currency symbol then convert to equivalent AED amount (eg. $10k becomes 36726.50, $1 becomes 3.67).\n\n"203        "Example:\n"204        "  For the query: \"properties near dubai mall under 43k\"\n"205        "  The expected output should be:\n"206        "    { \"surroundingArea\": \"dubai mall\", \"totalCosts\": 43000 }\n\n"207        "Return ONLY a valid JSON object with the extracted keys and their corresponding values, with no additional text."208    )209 210    human_str = f"Here is the query:\n{query}"211    filter_prompt = [212        {"role": "system", "content": system},213        {"role": "user", "content": human_str},214    ]215    response = llm_local.invoke(messages=filter_prompt)216    response_text = response.content if isinstance(response, AIMessage) else str(response)217    try:218        model_filters = extract_json_from_response(response_text)219    except Exception as e:220        print(f"JSON parsing error: {e}")221        model_filters = {}222    rule_filters = rule_based_extract(query)223    print("Rule-based extraction:", rule_filters)224    final_filters = {**model_filters, **rule_filters}225    print("Final extraction:", final_filters)226    return {"filters": final_filters}227 228 229@tool230def determine_route(query: str) -> dict:231    """For determining route using enhanced prompt and fallback logic."""232    # Define a set of keywords that are strong indicators of a real estate query.233    real_estate_keywords = estateKeywords234 235    # Check if the query includes any of the positive signals.236    pattern = re.compile("|".join(re.escape(keyword) for keyword in real_estate_keywords), re.IGNORECASE)237    positive_signal = bool(pattern.search(query))238 239    # Proceed with LLM classification regardless, but use the positive signal in fallback.240    llm_local = ChatQwen(temperature=0.3, streaming=False)241    transform_suggest_to_list = query.lower().replace("suggest ", "list ", -1)242    system = """243    Classify the user query as:244    245    - **"search"**: if it requests property listings with specific filters (e.g., location, price, property type like "2bhk", service charges, pet policies, etc.).246    - **"suggest"**: if it asks for property suggestions without filters.247    - **"detail"**: if it is asking for more information about a previously provided property (for example, "tell me more about property 5" or "I want more information regarding 4BHK").248    - **"general"**: for all other real estate-related questions.249    - **"out_of_domain"**: if the query is not related to real estate (for example, tourist attractions, restaurants, etc.).250    251    Keep in mind that queries mentioning terms like "service charge", "allow pets", "pet rules", etc., are considered real estate queries.252    253    Return only the keyword: search, suggest, detail, general, or out_of_domain.254    """255    human_str = f"Here is the query:\n{transform_suggest_to_list}"256    router_prompt = [257        {"role": "system", "content": system},258        {"role": "user", "content": human_str},259    ]260    261    response = llm_local.invoke(messages=router_prompt)262    response_text = response.content if isinstance(response, AIMessage) else str(response)263    route_value = str(response_text).strip().lower()264 265    # Fallback: if the query seems like a detailed request, override.266    detail_phrases = [267        "more information",268        "tell me more",269        "more details",270        "give me more details",271        "i need more details",272        "can you provide more details",273        "additional details",274        "further information",275        "expand on that",276        "explain further",277        "elaborate more",278        "more specifics",279        "i want to know more",280        "could you elaborate",281        "need more info",282        "provide more details",283        "detail it further",284        "in-depth information",285        "break it down further",286        "further explanation"287    ]288    if any(phrase in query.lower() for phrase in detail_phrases):289        route_value = "detail"290 291    if route_value not in {"search", "suggest", "detail", "general", "out_of_domain"}:292        route_value = "general"293    if route_value == "out_of_domain" and positive_signal:294        route_value = "general"295    if route_value == "out_of_domain":296        route_value = "general" if positive_signal else "out_of_domain"297        298    return {"route": route_value}299 300 301 302# ------------------------ Workflow Setup ------------------------303 304workflow = StateGraph(state_schema=dict)305 306def route_query(state: dict) -> dict:307    new_state = state.copy()308    try:309        new_state["route"] = determine_route.invoke(new_state.get("query", "")).get("route", "general")310        print(new_state["route"])311    except Exception as e:312        print(f"Routing error: {e}")313        new_state["route"] = "general"314    return new_state315 316def hybrid_extract(state: dict) -> dict:317    new_state = state.copy()318    new_state["filters"] = extract_filters.invoke(new_state.get("query", "")).get("filters", {})319    return new_state320 321def search_faiss(state: dict) -> dict:322    new_state = state.copy()323    query_embedding = st_model.encode([state["query"]])324    _, indices = index.search(query_embedding.astype(np.float32), 5)325    new_state["faiss_results"] = [docs[idx] for idx in indices[0] if idx < len(docs)]326    return new_state327 328def apply_filters(state: dict) -> dict:329    new_state = state.copy()330    new_state["final_results"] = apply_filters_partial(state["faiss_results"], state.get("filters", {}))331    return new_state332 333def suggest_properties(state: dict) -> dict:334    new_state = state.copy()335    new_state["suggestions"] = random.sample(docs, 5)336    return new_state337 338def handle_out_of_domain(state: dict) -> dict:339    new_state = state.copy()340    new_state["response"] = "I only handle real estate inquiries. Please ask a question related to properties."341    return new_state342 343 344def generate_response(state: dict) -> dict:345    new_state = state.copy()346    messages = []347    348    # Add the general query prompt.349    messages.append({"role": "system", "content": general_query_prompt})350    351    # If this is a detail query, add a system message that forces a detailed answer.352    if new_state.get("route", "general") == "detail":353        messages.append({354            "role": "system",355            "content": (356                "This is a detail query. Please provide detailed information about the property below. "357                "Do not generate a new list of properties; only use the provided property details to answer the query. "358                "Focus on answering the specific question (for example, whether pets are allowed)."359            )360        })361        362    # If property details are available, add them without clearing context.363    if new_state.get("current_properties"):364        property_context = format_property_data(new_state["current_properties"])365        messages.append({"role": "system", "content": "Available Property:\n" + property_context})366        # Do NOT clear current_properties here.367        messages.append({"role": "system", "content": "When responding, use only the provided property details to answer the user's specific question about the property."})368    369    # Add the conversation history.370    for msg in state.get("messages", []):371        if msg["role"] == "user":372            messages.append({"role": "user", "content": msg["content"]})373        else:374            messages.append({"role": "assistant", "content": msg["content"]})375    376    # Invoke the LLM with the constructed messages.377    connection_id = state.get("connection_id")378    loop = state.get("loop")379    if connection_id and loop:380        print("Yes")381        callback_manager = [WebSocketStreamingCallbackHandler(connection_id, loop)]382        _ = llm.invoke(383            messages,384            config={"callbacks": callback_manager}385        )386        new_state["response"] = ""387    else:388        callback_manager = [StreamingStdOutCallbackHandler()]389        response = llm.invoke(390            messages,391            config={"callbacks": callback_manager}392        )393        new_state["response"] = response.content if isinstance(response, AIMessage) else str(response)394    395    return new_state396 397def format_final_response(state: dict) -> dict:398    new_state = state.copy()399    # Only override the current_properties if this is NOT a detail query.400    if not state.get("route", "general") == "detail":401        if state.get("route") in ["search", "suggest"]:402            if "final_results" in state:403                new_state["current_properties"] = state["final_results"]404            elif "suggestions" in state:405                new_state["current_properties"] = state["suggestions"]406    407    # Then format the response based on the (possibly filtered) current_properties.408    if new_state.get("current_properties"):409        formatted = []410        for idx, prop in enumerate(new_state["current_properties"], 1):411            cost = prop.get("totalCosts", "N/A")412            cost_str = f"{cost:,}" if isinstance(cost, (int, float)) else cost413            formatted.append(414                f"{idx}. Type: {prop['propertyType']}, Cost: AED {cost_str}, "415                f"Size: {prop.get('propertySize', 'N/A')}, Amenities: {', '.join(map(str, prop.get('amenities', []))) if prop.get('amenities') else 'N/A'}, "416                f"Rental Yield: {prop.get('expectedRentalYield', 'N/A')}, "417                f"Ownership: {prop.get('ownershipType', 'N/A')}\n"418            )419        aggregated_response = "Here are the property details:\n" + "\n".join(formatted)420        connection_id = state.get("connection_id")421        loop = state.get("loop")422        if connection_id and loop:423            import time424            tokens = aggregated_response.split(" ")425            for token in tokens:426                asyncio.run_coroutine_threadsafe(427                    manager_socket.send_message(connection_id, token + " "),428                    loop429                )430                time.sleep(0.05)431            new_state["response"] = ""432        else:433            new_state["response"] = aggregated_response434    elif "response" in new_state:435        new_state["response"] = str(new_state["response"])436    return new_state437 438 439 440nodes = [441    ("route_query", route_query),442    ("hybrid_extract", hybrid_extract),443    ("faiss_search", search_faiss),444    ("apply_filters", apply_filters),445    ("suggest_properties", suggest_properties),446    ("handle_out_of_domain", handle_out_of_domain),447    ("generate_response", generate_response),448    ("format_response", format_final_response)449]450 451for name, node in nodes:452    workflow.add_node(name, node)453 454workflow.add_edge(START, "route_query")455workflow.add_conditional_edges(456    "route_query",457    lambda state: state.get("route", "general"),458    {459        "search": "hybrid_extract", 460        "suggest": "suggest_properties", 461        "detail": "generate_response", 462        "general": "generate_response", 463        "out_of_domain": "handle_out_of_domain"464    }465)466workflow.add_edge("hybrid_extract", "faiss_search")467workflow.add_edge("faiss_search", "apply_filters")468workflow.add_edge("apply_filters", "format_response")469workflow.add_edge("suggest_properties", "format_response")470workflow.add_edge("generate_response", "format_response")471workflow.add_edge("handle_out_of_domain", "format_response")472workflow.add_edge("format_response", END)473 474workflow_app = workflow.compile()475 476# ------------------------ Conversation Manager ------------------------477 478class ConversationManager:479    def __init__(self):480        self.conversation_history = []481        self.current_properties = []482 483    def _add_message(self, role: str, content: str):484        self.conversation_history.append({485            "role": role,486            "content": content,487            "timestamp": datetime.now().isoformat()488        })489 490    def process_query(self, query: str) -> str:491        # Reset context on greetings to avoid using off-domain history492        if query.strip().lower() in {"hi", "hello", "hey"}:493            self.conversation_history = []494            self.current_properties = []495            greeting_response = "Hello! How can I assist you today with your real estate inquiries?"496            self._add_message("assistant", greeting_response)497            return greeting_response498 499        try:500            self._add_message("user", query)501            initial_state = {502                "messages": self.conversation_history.copy(),503                "query": query,504                "route": "general",505                "filters": {},506                "current_properties": self.current_properties507            }508            for event in workflow_app.stream(initial_state, stream_mode="values"):509                final_state = event510            if 'final_results' in final_state:511                self.current_properties = final_state['final_results']512            elif 'suggestions' in final_state:513                self.current_properties = final_state['suggestions']514            if final_state.get("route") == "general":515                response_text = final_state.get("response", "")516                self._add_message("assistant", response_text)517                return response_text518            else:519                response = final_state.get("response", "I couldn't process that request.")520                self._add_message("assistant", response)521                return response522        except Exception as e:523            print(f"Processing error: {e}")524            return "Sorry, I encountered an error processing your request."525 526conversation_managers = {}527 528# ------------------------ FastAPI Backend with WebSockets ------------------------529 530app = FastAPI()531 532class ConnectionManager:533    def __init__(self):534        self.active_connections = {}535 536    async def connect(self, websocket: WebSocket):537        await websocket.accept()538        connection_id = str(uuid.uuid4())539        self.active_connections[connection_id] = websocket540        print(f"New connection: {connection_id}")541        return connection_id542 543    def disconnect(self, connection_id: str):544        if connection_id in self.active_connections:545            del self.active_connections[connection_id]546            print(f"Disconnected: {connection_id}")547 548    async def send_message(self, connection_id: str, message: str):549        websocket = self.active_connections.get(connection_id)550        if websocket:551            await websocket.send_text(message)552 553manager_socket = ConnectionManager()554 555def stream_query(query: str, connection_id: str, loop):556    conv_manager = conversation_managers.get(connection_id)557    if conv_manager is None:558        print(f"No conversation manager found for connection {connection_id}")559        return560 561    # Check for greetings and handle them immediately562    if query.strip().lower() in {"hi", "hello", "hey"}:563        conv_manager.conversation_history = []564        conv_manager.current_properties = []565        greeting_response = "Hello! How can I assist you today with your real estate inquiries?"566        conv_manager._add_message("assistant", greeting_response)567        asyncio.run_coroutine_threadsafe(568            manager_socket.send_message(connection_id, greeting_response),569            loop570        )571        return572 573    conv_manager._add_message("user", query)574    initial_state = {575        "messages": conv_manager.conversation_history.copy(),576        "query": query,577        "route": "general",578        "filters": {},579        "current_properties": conv_manager.current_properties,580        "connection_id": connection_id,581        "loop": loop582    }583    try:584        workflow_app.invoke(initial_state)585    except Exception as e:586        error_msg = f"Error processing query: {str(e)}"587        asyncio.run_coroutine_threadsafe(588            manager_socket.send_message(connection_id, error_msg),589            loop590        )591 592@app.websocket("/ws")593async def websocket_endpoint(websocket: WebSocket):594    connection_id = await manager_socket.connect(websocket)595    conversation_managers[connection_id] = ConversationManager()596    try:597        while True:598            query = await websocket.receive_text()599            loop = asyncio.get_event_loop()600            # loop = asyncio.get_running_loop()601            threading.Thread(602                target=stream_query,603                args=(query, connection_id, loop),604                daemon=True605            ).start()606    except WebSocketDisconnect:607        conv_manager = conversation_managers.get(connection_id)608        if conv_manager:609            filename = f"conversations/conversation_{connection_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"610            with open(filename, "w") as f:611                json.dump(conv_manager.conversation_history, f, indent=4)612            del conversation_managers[connection_id]613        manager_socket.disconnect(connection_id)614 615@app.post("/query")616async def post_query(query: str):617    conv_manager = ConversationManager()618    response = conv_manager.process_query(query)619    return {"response": response}620