import os
os.environ["OPENAI_API_KEY"] = "Insert OpenAI Key"

from langfuse.client import Langfuse

langfuse = Langfuse(
    public_key="<Insert PK>", 
    secret_key="<Insert SK>",
    host="Insert Host Details",
    release="1.0.0",
    debug=True
)
 
trace = langfuse.trace(
        name="LangChain-Agent-Trace", 
        user_id="user-1234"
)
 
langfuse_handler = trace.get_langchain_handler()

from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import tool

class SearchInput(BaseModel):
    rephrase_question: str = Field(description="Rephrased question for better context retrieval from Wiki knowledge base")

@tool
def should_retrieve(search_input: SearchInput) -> str:
    """
        Tool to retrieve the relevant context based on the rephrased question.
        When user query is about any Wiki related Knowledge Base context
    """
    return "retrieved" + search_input.rephrase_question

class ChatGPTInput(BaseModel):
    user_query: str = Field(description="Same user query passed in the input when question is not related to Wiki knowledge base or answer is available in the chat history itself.")

@tool
def use_like_chat_gpt(chat_gpt_input: ChatGPTInput) -> str:
    """
        Tool to use the language model like chat gpt when user query is not related to Wiki knowledge base or answer is available in the chat history itself.
    """
    return "retrieved" + chat_gpt_input.user_query

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)

from langchain_core.messages import SystemMessage

system_prompt = SystemMessage(
    content=
    '''
        You are an agent to determine if the given human question is about Wiki KB. 
        If user question is related to Wiki knowledge base then select should_retrieve tool and rephrase the human question for better context retrieval based on the given chat history and user query.
        If user question is not related to Wiki knowledge base then select use me like chat gpt tool.
    '''
)

from langchain_core.messages import HumanMessage, AIMessage
chat_history=[
        HumanMessage(content="Hello, who are you?"),
        AIMessage(content="I am doing great. How can I help you today?"),
        HumanMessage(content="I am looking for a restaurant in the center of town."),
        AIMessage(content="Ok, what is your budget?"),
        HumanMessage(content="I am on a budget of 30 dollars."),
        AIMessage(content="I have a few options for you: 1. Little Italy. 2. The Spot. 3. Corner Cafe."),
        HumanMessage(content="What type of food does each restaurant serve?"),
        AIMessage(content="Little Italy is an Italian restaurant. The Spot is a French restaurant. Corner Cafe is a coffee shop.")
    ]

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

prompt = ChatPromptTemplate.from_messages(
    [   system_prompt,
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ]
)

tools = [should_retrieve, use_like_chat_gpt]

from langchain.agents import AgentExecutor, create_openai_tools_agent

agent = create_openai_tools_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
agent_executor.invoke(
    {"input": "Who is the PM of India?", 
     "chat_history": chat_history
     }, {"callbacks":[langfuse_handler]}
)