import openai
import os
import dotenv
from llama_index.agent.openai import OpenAIAgent
from llama_index.llms.azure_openai import AzureOpenAI
from llama_index.core.tools.tool_spec.load_and_search.base import LoadAndSearchToolSpec
from llama_index.tools.google import GoogleSearchToolSpec
from llama_index.tools.weather import OpenWeatherMapToolSpec
from llama_index.embeddings.langchain import LangchainEmbedding
from llama_index.core import ServiceContext, set_global_service_context
from langchain.embeddings import OpenAIEmbeddings

# Get API keys from environment variables
dotenv.load_dotenv()
google_api_key = os.environ.get("GOOGLE_API_KEY")
google_cse_id = os.environ.get("GOOGLE_CSE_ID")
openweather_api_key = os.environ.get("OPENWEATHER_API_KEY")

os.environ["OPENAI_API_KEY"] = os.environ.get("AZURE_API_KEY")
openai.api_type = "azure"
openai.api_base = os.environ.get("AZURE_API_BASE")
openai.api_key = os.environ.get("AZURE_API_KEY")

azure_api_key = os.environ["AZURE_API_KEY"]
azure_api_type = "azure"
azure_api_base = os.environ.get("AZURE_API_BASE")
azure_api_version = os.environ.get("AZURE_API_VERSION")
azure_chatapi_version = os.environ.get("AZURE_CHATAPI_VERSION")

# Check if user set the davinci model flag
LLM_DEPLOYMENT_NAME = "gpt-4"
LLM_MODEL_NAME = "gpt-4"
EMBEDDINGS_DEPLOYMENT_NAME = "text-embedding-ada-002"
openai.api_version = os.environ.get("AZURE_CHATAPI_FUNCTIONS_VERSION")
print("Using gpt-4 2023-07-01-preview model.")

llm = AzureOpenAI(
    engine=LLM_DEPLOYMENT_NAME, 
    model=LLM_MODEL_NAME,
    api_key=azure_api_key,
    api_base=azure_api_base,
    api_type=azure_api_type,
    api_version=azure_chatapi_version,
    temperature=0.5,
    max_tokens=1024,
)

embedding_llm = LangchainEmbedding(
    OpenAIEmbeddings(
        model=EMBEDDINGS_DEPLOYMENT_NAME,
        deployment=EMBEDDINGS_DEPLOYMENT_NAME,
        openai_api_key=azure_api_key,
        openai_api_base=azure_api_base,
        openai_api_type=azure_api_type,
        openai_api_version=azure_api_version,
        chunk_size=16,
        max_retries=3,
    ),
    embed_batch_size=1,
)
service_context = ServiceContext.from_defaults(
    llm=llm,
    embed_model=embedding_llm,
    chunk_size=256,
)
set_global_service_context(service_context)

weather_tool = OpenWeatherMapToolSpec(
    key=openweather_api_key,
)

agent = OpenAIAgent.from_tools(
    weather_tool.to_tool_list(),
    llm=llm,
    verbose=True,
)

print("-------------------------------------------------------")
print(agent.chat("what is the current weather like in Hillsboro, Oregon"))

google_tool = GoogleSearchToolSpec(
    key=google_api_key,
    engine=google_cse_id
)

#Wrap the google search tool as it returns large payloads
tools = LoadAndSearchToolSpec.from_defaults(
    google_tool.to_tool_list()[0],
).to_tool_list()

agent = OpenAIAgent.from_tools(
    tools, 
    llm=llm,
    embedding_llm=embedding_llm,
    verbose=True
)
# agent = OpenAIAgent.from_tools(
#     google_tool.to_tool_list(), 
#     llm=llm, 
#     verbose=True
# )


print("-------------------------------------------------------")
print(agent.chat("whats the latest news summary about India?"))
print("-------------------------------------------------------")
print(agent.chat('who is Ange Postecoglou, where is he now and what is his philosophy about football?'))
print("-------------------------------------------------------")
print(agent.chat("whats the latest financial news summary?"))


langchain.embeddings.OpenAIEmbeddings
[((577, 597), 'dotenv.load_dotenv', 'dotenv.load_dotenv', ([], {}), '()\n', (595, 597), False, 'import dotenv\n'), ((615, 647), 'os.environ.get', 'os.environ.get', (['"""GOOGLE_API_KEY"""'], {}), "('GOOGLE_API_KEY')\n", (629, 647), False, 'import os\n'), ((664, 695), 'os.environ.get', 'os.environ.get', (['"""GOOGLE_CSE_ID"""'], {}), "('GOOGLE_CSE_ID')\n", (678, 695), False, 'import os\n'), ((718, 755), 'os.environ.get', 'os.environ.get', (['"""OPENWEATHER_API_KEY"""'], {}), "('OPENWEATHER_API_KEY')\n", (732, 755), False, 'import os\n'), ((788, 819), 'os.environ.get', 'os.environ.get', (['"""AZURE_API_KEY"""'], {}), "('AZURE_API_KEY')\n", (802, 819), False, 'import os\n'), ((864, 896), 'os.environ.get', 'os.environ.get', (['"""AZURE_API_BASE"""'], {}), "('AZURE_API_BASE')\n", (878, 896), False, 'import os\n'), ((914, 945), 'os.environ.get', 'os.environ.get', (['"""AZURE_API_KEY"""'], {}), "('AZURE_API_KEY')\n", (928, 945), False, 'import os\n'), ((1033, 1065), 'os.environ.get', 'os.environ.get', (['"""AZURE_API_BASE"""'], {}), "('AZURE_API_BASE')\n", (1047, 1065), False, 'import os\n'), ((1086, 1121), 'os.environ.get', 'os.environ.get', (['"""AZURE_API_VERSION"""'], {}), "('AZURE_API_VERSION')\n", (1100, 1121), False, 'import os\n'), ((1146, 1185), 'os.environ.get', 'os.environ.get', (['"""AZURE_CHATAPI_VERSION"""'], {}), "('AZURE_CHATAPI_VERSION')\n", (1160, 1185), False, 'import os\n'), ((1360, 1409), 'os.environ.get', 'os.environ.get', (['"""AZURE_CHATAPI_FUNCTIONS_VERSION"""'], {}), "('AZURE_CHATAPI_FUNCTIONS_VERSION')\n", (1374, 1409), False, 'import os\n'), ((1464, 1676), 'llama_index.llms.azure_openai.AzureOpenAI', 'AzureOpenAI', ([], {'engine': 'LLM_DEPLOYMENT_NAME', 'model': 'LLM_MODEL_NAME', 'api_key': 'azure_api_key', 'api_base': 'azure_api_base', 'api_type': 'azure_api_type', 'api_version': 'azure_chatapi_version', 'temperature': '(0.5)', 'max_tokens': '(1024)'}), '(engine=LLM_DEPLOYMENT_NAME, model=LLM_MODEL_NAME, api_key=\n    azure_api_key, api_base=azure_api_base, api_type=azure_api_type,\n    api_version=azure_chatapi_version, temperature=0.5, max_tokens=1024)\n', (1475, 1676), False, 'from llama_index.llms.azure_openai import AzureOpenAI\n'), ((2113, 2198), 'llama_index.core.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm': 'llm', 'embed_model': 'embedding_llm', 'chunk_size': '(256)'}), '(llm=llm, embed_model=embedding_llm, chunk_size=256\n    )\n', (2141, 2198), False, 'from llama_index.core import ServiceContext, set_global_service_context\n'), ((2209, 2252), 'llama_index.core.set_global_service_context', 'set_global_service_context', (['service_context'], {}), '(service_context)\n', (2235, 2252), False, 'from llama_index.core import ServiceContext, set_global_service_context\n'), ((2269, 2316), 'llama_index.tools.weather.OpenWeatherMapToolSpec', 'OpenWeatherMapToolSpec', ([], {'key': 'openweather_api_key'}), '(key=openweather_api_key)\n', (2291, 2316), False, 'from llama_index.tools.weather import OpenWeatherMapToolSpec\n'), ((2579, 2641), 'llama_index.tools.google.GoogleSearchToolSpec', 'GoogleSearchToolSpec', ([], {'key': 'google_api_key', 'engine': 'google_cse_id'}), '(key=google_api_key, engine=google_cse_id)\n', (2599, 2641), False, 'from llama_index.tools.google import GoogleSearchToolSpec\n'), ((2817, 2903), 'llama_index.agent.openai.OpenAIAgent.from_tools', 'OpenAIAgent.from_tools', (['tools'], {'llm': 'llm', 'embedding_llm': 'embedding_llm', 'verbose': '(True)'}), '(tools, llm=llm, embedding_llm=embedding_llm, verbose\n    =True)\n', (2839, 2903), False, 'from llama_index.agent.openai import OpenAIAgent\n'), ((1745, 2009), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'model': 'EMBEDDINGS_DEPLOYMENT_NAME', 'deployment': 'EMBEDDINGS_DEPLOYMENT_NAME', 'openai_api_key': 'azure_api_key', 'openai_api_base': 'azure_api_base', 'openai_api_type': 'azure_api_type', 'openai_api_version': 'azure_api_version', 'chunk_size': '(16)', 'max_retries': '(3)'}), '(model=EMBEDDINGS_DEPLOYMENT_NAME, deployment=\n    EMBEDDINGS_DEPLOYMENT_NAME, openai_api_key=azure_api_key,\n    openai_api_base=azure_api_base, openai_api_type=azure_api_type,\n    openai_api_version=azure_api_version, chunk_size=16, max_retries=3)\n', (1761, 2009), False, 'from langchain.embeddings import OpenAIEmbeddings\n')]