!pip install -q google-adk litellm python-dotenv exa-py tavily-python firecrawl-py

# Import necessary libraries for Google ADK, LLMs, and data processing
from google.adk.models.lite_llm import LiteLlm
from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner

# Import libraries for data retrieval and web scraping
from exa_py import Exa
from tavily import TavilyClient
from firecrawl import FirecrawlApp

# Import standard libraries for system operations and time handling
import os
from datetime import datetime, timedelta
from google.genai import types

os.environ["NEBIUS_API_BASE"] = "https://api.studio.nebius.ai/v1"
os.environ["NEBIUS_API_KEY"] = "your-nebius-api-key"
os.environ["EXA_API_KEY"] = "your-exa-api-key"
os.environ["TAVILY_API_KEY"] = "your-tavily-api-key"
os.environ["FIRECRAWL_API_KEY"] = "your-firecrawl-api-key"

nebius_model = LiteLlm(
    model="openai/meta-llama/Meta-Llama-3.1-8B-Instruct",
    api_base=os.getenv("NEBIUS_API_BASE"),
    api_key=os.getenv("NEBIUS_API_KEY")
)

# Tool 1: Define a function to search for AI-related news using Exa on X
def exa_search_ai(_: str) -> dict:
    results = Exa(api_key=os.getenv("EXA_API_KEY")).search_and_contents(
        query="Latest AI news OR new LLM models OR AI/Agents advancements",
        include_domains=["twitter.com", "x.com"],
        num_results=10,
        text=True,
        type="auto",
        highlights={"highlights_per_url": 2, "num_sentences": 3},
        start_published_date=(datetime.now() - timedelta(days=30)).isoformat()
    )
    return {"type": "exa", "results": [r.__dict__ for r in results.results]}


# Tool 2: Define a function to search artificialanalysis.ai using Tavily
def tavily_search_ai_analysis(_: str) -> dict:
    client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
    response = client.search(
        query="AI benchmarks OR AI/LLM statistics OR AI providers analysis",
        search_depth="advanced",
        time_range="week",
        include_domains=["artificialanalysis.ai"]
    )
    return {"type": "tavily", "results": response.get("results", [])}

# Tool 3: Define a function to Scrape content of Nebius AI using Firecrawl
def firecrawl_scrape_nebius(_: str) -> dict:
    firecrawl = FirecrawlApp(api_key=os.getenv("FIRECRAWL_API_KEY"))
    try:
        scrape_result = firecrawl.scrape_url(
            url="https://studio.nebius.com/",
            formats=["markdown"],
            only_main_content=True
        )
        if scrape_result.success:
            return {
                "type": "firecrawl",
                "markdown": scrape_result.markdown  # Directly access
            }
        else:
            return {
                "type": "firecrawl",
                "error": "Scraping failed."
            }
    except Exception as e:
        return {
            "type": "firecrawl",
            "error": str(e)
        }

# Agent 1: Define the Exa agent to fetch latest AI updates from X
exa_agent = LlmAgent(
    name="ExaAgent",
    model=nebius_model,
    description="Fetches latest AI news, LLMs, and advancements using Exa.",
    instruction="""
Use the exa_search_ai tool to fetch the latest information about AI, new LLMs, and advancements in the field from Twitter and X.
Prefix your response with " **🔥ExaAgent:** "
""",
    tools=[exa_search_ai],
    output_key="exa_news"
)

# Agent 2: Define the Tavily agent to fetch AI benchmarks and related infos
tavily_agent = LlmAgent(
    name="TavilyAgent",
    model=nebius_model,
    description="Fetches AI benchmarks, statistics, and analysis using Tavily.",
    instruction="""
Use the tavily_search_ai_analysis tool to retrieve benchmarks, statistics, llm providers and relevant analysis on AI.
Prefix your response with " **🐳TavilyAgent:** "
""",
    tools=[tavily_search_ai_analysis],
    output_key="tavily_news"
)

# Agent 3: Define the Summary agent to combine and format results from Exa and Tavily Agent
summary_agent = LlmAgent(
    name="SummaryAgent",
    model=nebius_model,
    description="Summarizes and formats Exa and Tavily results.",
    instruction="""
You are a summarizer and formatter.
- Combine the information from 'exa_results' and 'tavily_results'.
- Use markdown, bullet points, emojis (🚀, 📊, 📈).
- Prefix with " **🍥SummaryAgent:** "
""",
    tools=[],
    output_key="final_summary"
)

# Agent 4: Define the Firecrawl scrapper agent to scrape content of Nebius AI Studio
firecrawl_agent = LlmAgent(
    name="FirecrawlAgent",
    model=nebius_model,
    description="Scrapes Nebius Studio homepage using Firecrawl.",
    instruction="""
Use the firecrawl_scrape_nebius tool to fetch markdown content from Nebius Studio website in proper format.
Prefix your response with " **🔥FirecrawlAgent:** "
""",
    tools=[firecrawl_scrape_nebius],
    output_key="firecrawl_content"
)

# Agent 5: Define the Analysis agent to analyze the summary and provide insights and also consider Nebius Scrapped data
analysis_agent = LlmAgent(
    name="AnalysisAgent",
    model=LiteLlm(
        model="openai/nvidia/Llama-3_1-Nemotron-Ultra-253B-v1",
        api_base=os.getenv("NEBIUS_API_BASE"),
        api_key=os.getenv("NEBIUS_API_KEY")
    ),
    instruction="""
You are an AI analyst specializing in the latest AI trends and Large Language Models (LLMs).
- Analyze the 'final_summary', combining it with your knowledge of AI advancements and the information extracted from 'exa_results' and 'tavily_results'.
- Identify key trends, growth areas, and notable statistics related to AI and LLMs.
- Carefully examine the 'firecrawl_content', which contain data from Nebius AI Studio. This file provides details about available models on Nebius, including their names, pricing, token limits, and availability.
- Instead of focusing solely on model names, analyze the functional capabilities and intended use cases of LLMs mentioned in the 'final_summary'.
- Cross-reference the LLMs' functionalities with the Nebius AI Studio offerings in 'firecrawl_content', prioritizing models with similar features such as context window size, training data, or specialized capabilities.
- Utilize a markdown table to present a concise summary of these trends, focusing on specific LLMs and their providers available on Nebius AI Studio.
- Highlight context window supported by each LLM and pricing. This will help with choosing appropriate models based on specific use cases.

Prefix your response with "**🧠AnalysisAgent:**"
Generate the table and analysis at the end.
""",
    tools=[],
    output_key="ai_insights"
)

# Create the sequential agent and pass all agents in the correct order
sequential_agent = SequentialAgent(
    name="AITrendAnalysisAgent",
    agents=[exa_agent, tavily_agent, summary_agent, firecrawl_agent, analysis_agent]
)

# Define an in-memory session service to manage agent state
session_service = InMemorySessionService()

# Create a runner to execute the sequential agent
runner = Runner(agent=sequential_agent, session_service=session_service)

# Run the agent with a specific input
result = runner.run(inputs={"input": "Analyze the latest trends in AI."})

# Print the final output from the analysis agent
print(result["ai_insights"])