Initial commit with agent code and project structure

multi_tool_agent/agent.py

@@ -1,69 +1,136 @@
-from google.adk.agents import LlmAgent
-from google.adk.tools import google_search
+from google.adk.agents import Agent, LlmAgent
 from firecrawl import FirecrawlApp
-import google.genai as genai
 from dotenv import load_dotenv
-import os
-import json
-import time
+from pydantic import BaseModel
+from typing import List
+import os, time, re
 
-# Load environment variables
+# 🔹 Load API key from .env
 load_dotenv()
 FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY")
-
-# Setup Firecrawl
 firecrawl_app = FirecrawlApp(api_key=FIRECRAWL_API_KEY)
 
-# Setup Gemini model
-genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-gemini_model = genai.GenerativeModel("gemini-2.0-pro")
+# 🔹 Search Result Schema
+class SearchResult(BaseModel):
+    urls: List[str]
+
+# 🔹 Search Agent
+search_agent = LlmAgent(
+    name="SearchAgent",
+    model="gemini-2.0-flash",
+    instruction="""
+    You are a helpful web assistant. Given a travel prompt, return exactly 5 valid URLs relevant to that query.
 
-# Define a custom function to scrape and combine content
-def scrape_multiple_sites(urls):
-    combined_content = ""
+    Only return this format:
+    {"urls": ["https://site1.com", "https://site2.com", ...]}
+
+    Do not include any description or explanation.
+    """,
+    output_schema=SearchResult,
+    output_key="search_urls"
+)
+
+# 🔹 Scrape and Summarize Logic
+def scrape_and_summarize(user_prompt: str, urls: List[str]) -> str:
+    if not user_prompt or not urls:
+        return "Missing user_prompt or urls."
+
+    combined = ""
     for url in urls:
         try:
-            response = firecrawl_app.scrape_url(
-                url=url,
-                formats=["markdown"],
-                only_main_content=True,
-                timeout=120000
-            )
-            if hasattr(response, 'markdown') and response.markdown:
-                combined_content += f"\n\n---\n\n# {url}\n\n" + response.markdown
-            time.sleep(2)  # to respect rate limits
+            res = firecrawl_app.scrape_url(url=url, formats=["markdown"], only_main_content=True)
+            if hasattr(res, "markdown") and res.markdown:
+                combined += f"\n\n---\n# {url}\n\n{res.markdown}"
+            else:
+                print(f"No markdown content found at: {url}")
+            time.sleep(1.5)
         except Exception as e:
-            print(f"Failed to scrape {url}: {str(e)}")
-    return combined_content
-
-# Define a custom tool to wrap this logic as a callable tool
-def plan_trip_with_scraping(user_prompt):
-    # Step 1: Use Google search to get relevant URLs
-    search_results = google_search(user_prompt)
-    urls = [r["url"] for r in search_results[:5]]
-
-    # Step 2: Scrape each URL using Firecrawl
-    scraped_content = scrape_multiple_sites(urls)
-
-    if not scraped_content:
-        return "Sorry, couldn't find enough information."
-
-    # Step 3: Use Gemini to generate a travel plan
-    response = gemini_model.generate_content([
-        "You are a travel planner. Based on the scraped content below, generate a helpful day-wise travel itinerary.",
-        f"User Prompt: {user_prompt}",
-        f"Scraped Content: {scraped_content}"
-    ])
-    return response.text
-
-# Final agent definition
-agent = LlmAgent(
-    name="travel_planner",
-    description="Suggests travel plans using live Google Search results and content scraped from the top sites.",
-    instruction="Use google_search to find top URLs. Use Firecrawl to scrape the pages. Combine and use Gemini to create travel plan.",
-    tools=[google_search, plan_trip_with_scraping],
-    model="gemini-2.0-flash"
+            print(f"Failed to scrape {url} - {e}")
+
+    if not combined.strip():
+        return "Scraping failed or returned no usable data."
+
+    # Remove links to simplify text for LLM
+    trimmed = re.sub(r'\[.*?\]\(.*?\)', '', combined)[:6000]
+
+    return f"User Prompt: {user_prompt}\n\nScraped Content:\n{trimmed}"
+
+# 🔹 Scraper Agent
+scraper_agent = Agent(
+    name="ScraperAgent",
+    model="gemini-2.0-flash",
+    instruction="""
+You are a strict itinerary planner. You must generate a structured 3-day travel itinerary using ONLY the scraped content provided. 
+
+❌ DO NOT:
+- Write introductions or summaries like "Let's plan your trip..."
+- Mention assumptions like "I’ll assume you’re interested in..."
+- Ask the user for more information
+- Add personal opinions or suggestions
+
+✅ DO:
+- Only use scraped content
+- Follow this exact format:
+
+Title: "3-Day Trip to <Destination> under ₹<Budget>"
+
+Day 1: <Title>
+Morning:
+- <Activity>
+Afternoon:
+- <Activity>
+Evening:
+- <Activity>
+Night:
+- <Activity>
+Estimated Day 1 Spend: ₹<Amount>
+
+Day 2: ...
+Day 3: ...
+
+Budget Breakdown:
+- Transport: ₹<amount>
+- Food: ₹<amount>
+- Accommodation: ₹<amount>
+- Sightseeing: ₹<amount>
+- Total: ₹<amount>
+
+Only generate this format. If content is missing, say: "Not enough data to generate a valid itinerary."
+"""
 )
 
-root_agent = agent
+
+# 🔹 Root Agent
+class TravelRootAgent(Agent):
+    def run(self, user_prompt: str):
+        
+        # STEP 1: Get 5 travel URLs
+        search_result = search_agent.run(user_prompt)
+        urls = search_result.get("urls", [])
+
+        for u in urls: print("  -", u)
+
+        if not urls:
+            return "Search failed — no URLs returned."
+
+        # STEP 2: Scrape and summarize
+        scraped_content = scrape_and_summarize(user_prompt, urls)
+        print("\n Scraped Content Preview:\n", scraped_content[:1000])
+
+        # STEP 3: Send to ScraperAgent
+        payload = f"""
+user_prompt: {user_prompt}
+
+scraped_content:
+{scraped_content}
+"""
+        
+        return scraper_agent.run(payload)
+
+# ✅ Instantiate root agent
+root_agent = TravelRootAgent(
+    name="TravelRootAgent",
+    model="gemini-2.0-flash",
+    instruction="Plan a 3-day travel itinerary using external scraped travel content."
+)