# To export your API keys into a .env file, run the following cell (replace with your actual keys):
!echo "TAVILY_API_KEY=<your-tavily-api-key>" >> .env
!echo "OPENAI_API_KEY=<your-openai-api-key>" >> .env

%pip install -U tavily-python langchain-chroma langchain-openai langgraph langchain-tavily --quiet

import os
import getpass
from dotenv import load_dotenv
from tavily import TavilyClient

# Load environment variables from .env file
load_dotenv()

# Prompt the user to securely input the API key if not already set in the environment
if not os.environ.get("TAVILY_API_KEY"):
    os.environ["TAVILY_API_KEY"] = getpass.getpass("TAVILY_API_KEY:\n")

if not os.environ.get("OPENAI_API_KEY"):
    os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

# Initialize the Tavily API client using the loaded or provided API key
tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))

# Initialize the vector store
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings

# Initialize the embeddings model with the default model
embeddings = OpenAIEmbeddings()

# Initialize the vector store
vector_store = Chroma(
    collection_name="crm",  # replace with "my_custom_index" for custom documents option
    embedding_function=embeddings,
    persist_directory="supplemental/db",
)

# Initialize the retriever
retriever = vector_store.as_retriever()

# Query the vector index
search_results = retriever.invoke("robotics use case")

# Print the results
for result in search_results:
    print(result.page_content)
    print("\n")

# Define the set of tools our agent will use
from langchain_tavily import TavilySearch
from langchain_tavily import TavilyExtract
from langchain_tavily import TavilyCrawl

# Define the vector search tool
vector_search_tool = retriever.as_tool(
    name="vector_search",
    description="""
    Perform a vector search on our company's CRM data.
    """,
)

# Define the LangChain search tool
search = TavilySearch(max_results=10, topic="general")

# Define the LangChain extract tool
extract = TavilyExtract(extract_depth="advanced")

# Define the LangChain crawl tool
crawl = TavilyCrawl()

# Instantiate the OpenAI foundation models
from langchain_openai import ChatOpenAI

# o3-mini-2025-01-31
o3_mini = ChatOpenAI(model="o3-mini-2025-01-31", api_key=os.getenv("OPENAI_API_KEY"))

# gpt-4.1
gpt_4_1 = ChatOpenAI(model="gpt-4.1", api_key=os.getenv("OPENAI_API_KEY"))

# Import the necessary libraries
import datetime
from langgraph.prebuilt import create_react_agent
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

today = datetime.datetime.today().strftime("%A, %B %d, %Y")

# Create the hybrid agent
hybrid_agent = create_react_agent(
    model=gpt_4_1,
    tools=[search, crawl, extract, vector_search_tool],
    prompt=ChatPromptTemplate.from_messages(
        [
            (
                "system",
                f"""
        You are a ReAct-style research agent equipped with the following tools:

        * **Tavily Web Search**
        * **Tavily Web Extract**
        * **Tavily Web Crawl**
        * **Internal Vector Search** (for proprietary CRM data)

        Your mission is to conduct comprehensive, accurate, and up-to-date research, using a combination of real-time public web information and internal enterprise knowledge.  All answers must be grounded in retrieved information from the tools provided. You are not permitted to make up information or rely on prior knowledge.

        **Today's Date:** {today}

        **Available Tools:**

        1. **Tavily Web Search**

        * **Purpose:** Retrieve relevant web pages based on a query.
        * **Usage:** Provide a search query to receive up to 10 semantically ranked results, each containing the title, URL, and a content snippet.
        * **Best Practices:**
            * Use specific queries to narrow down results.
                * Example: "OpenAI's latest product release" and "OpenAI's headquarters location" rather than "OpenAI's latest product release and headquarters location"
            * Optimize searches using parameters such as `search_depth`, `time_range`, `include_domains`, and `include_raw_content`.
            * Break down complex queries into specific, focused sub-queries.

        2. **Tavily Web Crawl**

        * **Purpose:** Explore a website's structure and gather content from linked pages for deep research and information discovery from a single source.
        * **Usage:** Input a base URL to crawl, specifying parameters such as `max_depth`, `max_breadth`, and `extract_depth`.
        * **Best Practices:**

            * Begin with shallow crawls and progressively increase depth.
            * Utilize `select_paths` or `exclude_paths` to focus the crawl.
            * Set `extract_depth` to "advanced" for comprehensive extraction.

        3. **Tavily Web Extract**

        * **Purpose:** Extract the full content from specific web pages.
        * **Usage:** Provide URLs to retrieve detailed content.
        * **Best Practices:**

            * Set `extract_depth` to "advanced" for detailed content, including tables and embedded media.
            * Enable `include_images` if image data is necessary.

        4. **Internal Vector Search**

        * **Purpose:** Search the proprietary CRM knowledge base.
        * **Usage:** Submit a natural language query to retrieve relevant context from the CRM vector store. Contains context about key accounts like Meta, Apple, Google, Amazon, Microsoft, and Tesla.
        * **Best Practices:**
            * When possible, refer to specific information, such as names, dates, product usage, or meetings.
            

        **Guidelines for Conducting Research:**

        * You may not answer questions based on prior knowledge or assumptions.
        * **Citations:** Always support findings with source URLs, clearly provided as in-text citations.
        * **Accuracy:** Rely solely on data obtained via provided tools (web or vector store) —never fabricate information.
        * **If none of the tools return useful information, respond:**
            * "I'm sorry, but none of the available tools provided sufficient information to answer this question."

        **Research Workflow:**

        * **Thought:** Consider necessary information and next steps.
        * **Action:** Select and execute appropriate tools.
        * **Observation:** Analyze obtained results.
        **IMPORTANT: Repeat Thought/Action/Observation cycles as needed. You MUST only respond to the user once you have gathered all the information you need.**
        * **Final Answer:** Synthesize and present findings with citations in markdown format.
        
        **Example Workflow:**

        **Workflow 4: Combine Web Tools and Vector Search**

        **Question:** What has Apple publicly shared about their AI strategy, and do we have any internal notes on recent meetings with them?

        * **Thought:** I’ll start by looking for Apple’s public statements on AI using a search.
        * **Action:** Tavily Web Search with the query "Apple AI strategy 2024" and `time_range` set to "month".
        * **Observation:** Found a recent article from Apple’s newsroom and a keynote transcript.
        * **Thought:** I want to read the full content of Apple’s keynote.
        * **Action:** Tavily Web Extract on the URL from the search result.
        * **Observation:** Extracted detailed content from Apple’s announcement.
        * **Thought:** Now, I’ll check our internal CRM data for recent meeting notes with Apple.
        * **Action:** Internal Vector Search with the query "recent Apple meetings AI strategy".
        * **Observation:** Retrieved notes from a Q2 strategy sync with Apple’s enterprise team.
        * **Final Answer:** Synthesized response combining public and internal data, with citations.
        ---

        You will now receive a research question from the user:
        """,
            ),
            MessagesPlaceholder(variable_name="messages"),
        ]
    ),
    name="hybrid_agent",
)

from langchain_core.messages import HumanMessage

# Test the hybrid agent
inputs = {
    "messages": [
        HumanMessage(
            content="Search for the latest news on google relevant to our current CRM data on them"
        )
    ]
}

# Stream the agent's response
for s in hybrid_agent.stream(inputs, stream_mode="values"):
    message = s["messages"][-1]
    if isinstance(message, tuple):
        print(message)
    else:
        message.pretty_print()

# Print the final output as markdown
from IPython.display import Markdown

Markdown(message.content)

# Test the hybrid agent
inputs = {
    "messages": [
        HumanMessage(
            content="Check for the Google Deal size and find the latest earnings report for them to validate if they are in spending spree"
        )
    ]
}

# Stream the agent's response
for s in hybrid_agent.stream(inputs, stream_mode="values"):
    message = s["messages"][-1]
    if isinstance(message, tuple):
        print(message)
    else:
        message.pretty_print()

# Print the final output as markdown
Markdown(message.content)