Update main.py

main.py

@@ -1,233 +1,131 @@
-# # main.py
-# import os
-# from fastapi import FastAPI, HTTPException
-# from fastapi.middleware.cors import CORSMiddleware
-# from dotenv import load_dotenv
-#
-# from models import TravelDataInput, PlanResponse, ChatInput, ChatResponse
-# from image_searcher import ImageSearcher, create_image_queries, extract_destination
-# from gemini_utils import create_travel_plan_prompt, generate_with_gemini, create_chat_prompt
-# from rag_utils import get_relevant_plan
-#
-# # Load environment variables (especially API keys)
-# load_dotenv()
-# PEXELS_API_KEY = os.getenv("PEXELS_API_KEY", "LB4mRviLcvE72R2645m3f4NpYGBMpPucVctpvzJAsCJYC2dhbGVpLraK") # Replace or set in env
-# GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "AIzaSyDuDPP0eWWA0UsGBBIeOI7Fl0WcPli4Sdo") # Replace or set in env
-#
-# if not PEXELS_API_KEY or not GOOGLE_API_KEY:
-#     print("Warning: API keys not found. Please set PEXELS_API_KEY and GOOGLE_API_KEY environment variables.")
-#
-# # Initialize app and image searcher
-# app = FastAPI(title="Travel Recommendation API")
-# image_searcher = ImageSearcher(PEXELS_API_KEY)
-#
-# # CORS (Cross-Origin Resource Sharing) - Allow frontend access
-# app.add_middleware(
-#     CORSMiddleware,
-#     allow_origins=["*"], # Adjust in production
-#     allow_credentials=True,
-#     allow_methods=["*"],
-#     allow_headers=["*"],
-# )
-#
-# @app.get("/", summary="Root endpoint for health check")
-# def read_root():
-#     return {"message": "Travel Recommendation API is running!"}
-#
-# @app.post("/generate_plan", response_model=PlanResponse, summary="Generate a travel plan")
-# async def generate_plan_endpoint(travel_input: TravelDataInput):
-#     travel_data = travel_input.model_dump() # Use model_dump() for Pydantic v2
-#     source = "generated"
-#     plan_text = None
-#
-#     # 1. Try RAG first
-#     try:
-#         relevant_plan = get_relevant_plan(travel_data)
-#         if relevant_plan:
-#             plan_text = relevant_plan
-#             source = "index"
-#     except Exception as e:
-#         print(f"RAG check failed (continuing with generation): {e}")
-#
-#     # 2. If no RAG plan, generate with Gemini
-#     if not plan_text:
-#         try:
-#             prompt = create_travel_plan_prompt(travel_data)
-#             plan_text = generate_with_gemini(prompt)
-#             source = "generated"
-#         except Exception as e:
-#             print(f"Gemini generation failed: {e}")
-#             raise HTTPException(status_code=500, detail=f"Failed to generate plan: {e}")
-#
-#     if not plan_text:
-#         raise HTTPException(status_code=500, detail="Could not retrieve or generate a travel plan.")
-#
-#     # 3. Extract destination & get images
-#     try:
-#         destination = extract_destination(plan_text, travel_data) or "Suggested Destination"
-#         image_queries = create_image_queries(travel_data, destination)
-#         images = image_searcher.search_all_sources(image_queries, num_images=6)
-#     except Exception as e:
-#         print(f"Image search failed: {e}")
-#         images = [] # Return empty list if image search fails
-#         image_queries = []
-#
-#     return PlanResponse(
-#         plan_text=plan_text,
-#         destination_used=destination,
-#         images=images,
-#         image_queries_used=image_queries,
-#         source=source
-#     )
-#
-# @app.post("/chat", response_model=ChatResponse, summary="Chat about the travel plan")
-# async def chat_endpoint(chat_input: ChatInput):
-#     try:
-#         prompt = create_chat_prompt(
-#             chat_input.plan_text,
-#             chat_input.destination,
-#             chat_input.num_days,
-#             [msg.model_dump() for msg in chat_input.chat_history], # Convert Pydantic to dict
-#             chat_input.user_input
-#         )
-#         response = generate_with_gemini(prompt, max_tokens=1024, temperature=0.5)
-#         return ChatResponse(assistant_response=response)
-#     except Exception as e:
-#         print(f"Chat generation failed: {e}")
-#         raise HTTPException(status_code=500, detail=f"Failed to get chat response: {e}")
-#
-# # To run locally: uvicorn main:app --reload
-
-# PEXELS_API_KEY = "LB4mRviLcvE72R2645m3f4NpYGBMpPucVctpvzJAsCJYC2dhbGVpLraK"
-# GOOGLE_API_KEY = "AIzaSyDuDPP0eWWA0UsGBBIeOI7Fl0WcPli4Sdo"
-
-import re
-import os
-import json
-from typing import List
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-
-from models import TravelDataInput, PlanResponse
-from image_searcher import ImageSearcher, create_image_queries, extract_destination
-from gemini_utils import create_travel_plan_prompt, generate_with_gemini
-from rag_utils import get_relevant_plan
-
-# API keys
-PEXELS_API_KEY = "LB4mRviLcvE72R2645m3f4NpYGBMpPucVctpvzJAsCJYC2dhbGVpLraK"
-GOOGLE_API_KEY = "AIzaSyDuDPP0eWWA0UsGBBIeOI7Fl0WcPli4Sdo"
-
-# App setup
-app = FastAPI(title="Travel Recommendation API")
-image_searcher = ImageSearcher(PEXELS_API_KEY)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-@app.get("/", summary="Root endpoint for health check")
-def read_root():
-    return {"message": "Travel Recommendation API is running!"}
-
-@app.post("/generate_plan", response_model=PlanResponse, summary="Generate a travel plan")
-async def generate_plan_endpoint(travel_input: TravelDataInput):
-    travel_data = travel_input.model_dump()
-    source = "generated"
-    plan_text = None
-
-    # Step 1: RAG
-    try:
-        relevant_plan = get_relevant_plan(travel_data)
-        if relevant_plan:
-            plan_text = relevant_plan
-            source = "index"
-    except Exception as e:
-        print(f"RAG check failed: {e}")
-
-    # Step 2: Generate if RAG fails
-    if not plan_text:
-        try:
-            prompt = create_travel_plan_prompt(travel_data)
-            plan_text = generate_with_gemini(prompt)
-        except Exception as e:
-            print(f"Gemini generation failed: {e}")
-            raise HTTPException(status_code=500, detail=f"Failed to generate plan: {e}")
-
-    if not plan_text:
-        raise HTTPException(status_code=500, detail="Could not retrieve or generate a travel plan.")
-
-    # Step 3: Get destination & images
-    try:
-        destination = extract_destination(plan_text, travel_data) or "Suggested Destination"
-        image_queries = create_image_queries(travel_data, destination)
-        images = image_searcher.search_all_sources(image_queries, num_images=6)
-    except Exception as e:
-        print(f"Image search failed: {e}")
-        images = []
-        image_queries = []
-
-    return PlanResponse(
-        plan_text=plan_text,
-        destination_used=destination,
-        images=images,
-        image_queries_used=image_queries,
-        source=source,
-    )
-
-# ========== EXTRACTION ENDPOINT ==========
-
-class ExtractionInput(BaseModel):
-    plan_text: str
-
-@app.post("/extract_package", summary="Extract structured package data from plan text")
-async def extract_package_endpoint(data: ExtractionInput):
-    extraction_prompt = f"""
-Extract ONLY the following fields from the travel plan below and return STRICTLY in raw JSON format. No extra explanation, no markdown, no text.
-
-Required JSON format:
-{{
-  "packageName": string,
-  "packageDescription": string,
-  "packageDestination": string,
-  "packageDays": integer,
-  "packageNights": integer,
-  "packageAccommodation": string,
-  "packageTransportation": string,
-  "packageMeals": string,
-  "packageActivities": string,
-  "packagePrice": float,
-  "packageDiscountPrice": float,
-  "packageOffer": boolean,
-  "packageRating": 0,
-  "packageTotalRatings": 0
-}}
-
-TRAVEL PLAN:
-\"\"\" 
-{data.plan_text}
-\"\"\"
-"""
-
-    try:
-        raw_response = generate_with_gemini(extraction_prompt)
-        print("Gemini raw response:", raw_response)
-
-        # Extract the JSON body from the response
-        json_match = re.search(r'\{.*\}', raw_response, re.DOTALL)
-        if not json_match:
-            raise ValueError("No JSON object found in Gemini response")
-
-        json_str = json_match.group(0)
-        structured_package = json.loads(json_str)
-
-    except Exception as e:
-        print("Gemini extraction failed:", e)
-        print("Prompt was:\n", extraction_prompt)
-        raise HTTPException(status_code=500, detail="Failed to extract structured package data.")
-
-    return structured_package
+from dotenv import load_dotenv
+import re
+import os
+import json
+from typing import List
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+from models import TravelDataInput, PlanResponse
+from image_searcher import ImageSearcher, create_image_queries, extract_destination
+from gemini_utils import create_travel_plan_prompt, generate_with_gemini
+from rag_utils import get_relevant_plan
+
+
+load_dotenv()
+PEXELS_API_KEY = os.getenv("PEXELS_API_KEY") 
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+# App setup
+app = FastAPI(title="Travel Recommendation API")
+image_searcher = ImageSearcher(PEXELS_API_KEY)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/", summary="Root endpoint for health check")
+def read_root():
+    return {"message": "Travel Recommendation API is running!"}
+
+@app.post("/generate_plan", response_model=PlanResponse, summary="Generate a travel plan")
+async def generate_plan_endpoint(travel_input: TravelDataInput):
+    travel_data = travel_input.model_dump()
+    source = "generated"
+    plan_text = None
+
+    # Step 1: RAG
+    try:
+        relevant_plan = get_relevant_plan(travel_data)
+        if relevant_plan:
+            plan_text = relevant_plan
+            source = "index"
+    except Exception as e:
+        print(f"RAG check failed: {e}")
+
+    # Step 2: Generate if RAG fails
+    if not plan_text:
+        try:
+            prompt = create_travel_plan_prompt(travel_data)
+            plan_text = generate_with_gemini(prompt)
+        except Exception as e:
+            print(f"Gemini generation failed: {e}")
+            raise HTTPException(status_code=500, detail=f"Failed to generate plan: {e}")
+
+    if not plan_text:
+        raise HTTPException(status_code=500, detail="Could not retrieve or generate a travel plan.")
+
+    # Step 3: Get destination & images
+    try:
+        destination = extract_destination(plan_text, travel_data) or "Suggested Destination"
+        image_queries = create_image_queries(travel_data, destination)
+        images = image_searcher.search_all_sources(image_queries, num_images=6)
+    except Exception as e:
+        print(f"Image search failed: {e}")
+        images = []
+        image_queries = []
+
+    return PlanResponse(
+        plan_text=plan_text,
+        destination_used=destination,
+        images=images,
+        image_queries_used=image_queries,
+        source=source,
+    )
+
+# ========== EXTRACTION ENDPOINT ==========
+
+class ExtractionInput(BaseModel):
+    plan_text: str
+
+@app.post("/extract_package", summary="Extract structured package data from plan text")
+async def extract_package_endpoint(data: ExtractionInput):
+    extraction_prompt = f"""
+Extract ONLY the following fields from the travel plan below and return STRICTLY in raw JSON format. No extra explanation, no markdown, no text.
+
+Required JSON format:
+{{
+  "packageName": string,
+  "packageDescription": string,
+  "packageDestination": string,
+  "packageDays": integer,
+  "packageNights": integer,
+  "packageAccommodation": string,
+  "packageTransportation": string,
+  "packageMeals": string,
+  "packageActivities": string,
+  "packagePrice": float,
+  "packageDiscountPrice": float,
+  "packageOffer": boolean,
+  "packageRating": 0,
+  "packageTotalRatings": 0
+}}
+
+TRAVEL PLAN:
+\"\"\" 
+{data.plan_text}
+\"\"\"
+"""
+
+    try:
+        raw_response = generate_with_gemini(extraction_prompt)
+        print("Gemini raw response:", raw_response)
+
+        # Extract the JSON body from the response
+        json_match = re.search(r'\{.*\}', raw_response, re.DOTALL)
+        if not json_match:
+            raise ValueError("No JSON object found in Gemini response")
+
+        json_str = json_match.group(0)
+        structured_package = json.loads(json_str)
+
+    except Exception as e:
+        print("Gemini extraction failed:", e)
+        print("Prompt was:\n", extraction_prompt)
+        raise HTTPException(status_code=500, detail="Failed to extract structured package data.")
+
+    return structured_package