# -*- coding: utf-8 -*-
"""Agentic_RAG_examples.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1ew19V8KYJPiNwBGaWBB704_QeLYR8IkM
"""

from google.colab import userdata
from typing import Literal
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langchain.chat_models import init_chat_model
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings # Updated import
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph, MessagesState
from langgraph.prebuilt import ToolNode
import os
import pypdf


@tool
# define a tool for context retrieval
def retrieve_context(query: str):
  """Search for relevant documents."""
  # parse and load the PDF document
  docs = []
  path = "/content/"

  for pdf in os.listdir(path):
    if pdf.endswith(".pdf"):
      loader = PyPDFLoader(path+pdf)
    docs.extend(loader.load())

  from langchain_text_splitters import CharacterTextSplitter
  text_splitter= CharacterTextSplitter(chunk_size=5, chunk_overlap=0)
  split_docs = text_splitter.split_documents(docs)

  # load the document into a knowledge base
  faiss_index = FAISS.from_documents(split_docs, OpenAIEmbeddings())
  faiss_index.save_local("File storage")

  # load and saved embeddings
  embeddings = OpenAIEmbeddings()
  loaded_vectors= FAISS.load_local("File storage", embeddings, allow_dangerous_deserialization=True)

  retriever = loaded_vectors.as_retriever()
  results = retriever.invoke(query)

  return "\n".join([doc.page_content for doc in results])


tools = [retrieve_context]
tool_node = ToolNode(tools)

# set the openai key and define the LLM
os.environ['OPENAI_API_KEY'] = userdata.get('openai')

model = ChatOpenAI(model="o4-mini-high").bind_tools(tools)

# Function to decide whether to continue or stop the workflow
def should_continue(state: MessagesState) -> Literal["tools", END]:
  messages = state['messages']
  last_message = messages[-1]
  # If the LLM makes a tool call, go to the "tools" node
  if last_message.tool_calls:
    return "tools"

  # Otherwise, finish the workflow
  return END

# Function that invokes the model
def call_model(state: MessagesState):

  messages = state['messages']
  response = model.invoke(messages)

  return {"messages": [response]}  # Returns as a list to add to the state

# Define the workflow with LangGraph
workflow = StateGraph(MessagesState)

# Add nodes to the graph
workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

# Connect nodes
workflow.add_edge(START, "agent") # Initial Entry
workflow.add_conditional_edges("agent", should_continue) # Decision after the "agent" node
workflow.add_edge("tools", "agent") # Cycle between tools and agent

# Configure memory to persist the state
Checkpointer = MemorySaver()

# Compile the graph into a LangChain Runnable application
app = workflow.compile(checkpointer = Checkpointer)


def rag(query):
  # Execute the workflow
  final_state = app.invoke(
    {"messages": [HumanMessage(content=query)]},
    config={"configurable": {"thread_id": 42}}
  )

  # show the final response
  return final_state["messages"][-1].content

rag("How to set up a Spinning bike?")