Update main.py

main.py

@@ -144,59 +144,38 @@ TRAVEL PLAN:
     return structured_package
 
 
+
 class ChatRequest(BaseModel):
     plan_text: str
     user_query: str
 
-    
-@app.post("/chat", summary="Answer user questions about the generated travel plan")
-async def chat_about_plan(chat_input: ChatRequest):
-    plan_text = chat_input.plan_text
-    user_query = chat_input.user_query.lower()
+def create_gemini_prompt(plan_text: str, user_query: str) -> str:
+    return f"""
+You are a helpful travel assistant specialized in tours and travel within Pakistan. 
 
-    if not is_query_related_to_travel_pakistan(user_query):
-        return {
-            "response": "❌ I can only answer questions related to the travel plan or tours within Pakistan."
-        }
+Here is a travel plan describing a trip:
 
-    try:
-        if "best time" in user_query or "when to visit" in user_query or "weather" in user_query:
-            match = re.search(r"\*\*Best Time to Visit:\*\*\s*(.*?)\n", plan_text)
-            return {
-                "response": match.group(1).strip() if match else "The best time to visit isn't mentioned in the plan."
-            }
-
-        if "cost" in user_query or "price" in user_query or "budget" in user_query:
-            match = re.search(r"\*\*Total Estimated Cost:\*\*\s*(.*?)\n", plan_text)
-            return {
-                "response": f"The estimated cost is {match.group(1).strip()}." if match else "The cost isn't specified in the plan."
-            }
-
-        if "activities" in user_query or "things to do" in user_query or "highlights" in user_query:
-            match = re.search(r"\*\*Key Highlights and Attractions:\*\*\s*(.*?)\n", plan_text)
-            return {
-                "response": match.group(1).strip() if match else "Activities and highlights are not clearly mentioned."
-            }
-
-        if "food" in user_query or "meal" in user_query or "eat" in user_query or "restaurant" in user_query:
-            # Extract lunch or dinner mentions
-            food_section = re.findall(r"\*\*Afternoon.*?\*\*Lunch.*?:\*\*\s*(.*?)\n|\*\*Evening.*?\*\*Dinner.*?:\*\*\s*(.*?)\n", plan_text, re.DOTALL)
-            flattened = [item for tup in food_section for item in tup if item]
-            return {
-                "response": "Here’s what you can eat: " + " ".join(flattened) if flattened else "Food recommendations aren't available."
-            }
-
-        if "transport" in user_query or "how to get" in user_query:
-            match = re.search(r"\*\*Taxila to Murree:\*\*\s*(.*?)\n", plan_text)
-            return {
-                "response": match.group(1).strip() if match else "Transportation details are not mentioned."
-            }
-
-        # Fallback with section list
-        return {
-            "response": "I can help with specific parts of the travel plan like best time to visit, cost, transport, food, and highlights. What would you like to know?"
-        }
+\"\"\"
+{plan_text}
+\"\"\"
+
+Please answer the user's question below **only if it is related to travel, sightseeing, food, transport, costs, culture, or activities in Pakistan based on the above plan**.
 
+User question:
+{user_query}
+
+If the question is unrelated to travel in Pakistan, respond politely with:
+"I can only answer questions related to travel and tours within Pakistan."
+
+Answer concisely and clearly.
+"""
+
+@app.post("/chat", summary="Answer user questions about the generated travel plan")
+async def chat_about_plan(chat_input: ChatRequest):
+    try:
+        prompt = create_gemini_prompt(chat_input.plan_text, chat_input.user_query)
+        llm_response = generate_with_gemini(prompt)  # Remove 'await' if your wrapper is not async
+        return {"response": llm_response.strip()}
     except Exception as e:
-        print("Chat processing failed:", e)
+        print("Error in chat endpoint:", e)
         raise HTTPException(status_code=500, detail="Failed to process your question.")