rag + websearch

Modules/rag.py

@@ -1,55 +1,68 @@
 import os
 os.environ['TOKENIZERS_PARALLELISM'] = 'true'
 os.environ['MISTRAL_API_KEY'] = "i5jSJkCFNGKfgIztloxTMjfckiFbYBj4"
-os.environ['OPENAI_API_KEY'] = ""
+# os.environ['OPENAI_API_KEY'] = "sk-proj-2WJfO8JpVyrdIeJ8QsO0T3BlbkFJWLhZF1xMlRZVFjNBccWh"
 os.environ['TAVILY_API_KEY'] = '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'
 
 mistral_api_key = os.getenv("MISTRAL_API_KEY")
 tavily_api_key = os.getenv("TAVILY_API_KEY")
-
+from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import WebBaseLoader
 from langchain_community.vectorstores import Chroma, FAISS
 from langchain_mistralai import MistralAIEmbeddings
-from langchain_openai import OpenAIEmbeddings
+from langchain import hub
 from typing import Literal
-
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_mistralai import ChatMistralAI
-from sentence_transformers import SentenceTransformer
 from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
-from transformers import AutoModel, AutoTokenizer
-from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+from langchain_community.tools import DuckDuckGoSearchRun
 
-urls = [
-    "https://lilianweng.github.io/posts/2023-06-23-agent/",
-    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
-    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
-]
+# urls = [
+#     "https://www.toutelanutrition.com/wikifit/guide-nutrition/nutrition-sportive/apports-proteines",
 
-docs = [WebBaseLoader(url).load() for url in urls]
-docs_list = [item for sublist in docs for item in sublist]
+# ]
 
-text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
-    chunk_size=250, chunk_overlap=0
-)
-doc_splits = text_splitter.split_documents(docs_list)
-
-##################### EMBED #####################
-# embeddings = MistralAIEmbeddings(mistral_api_key=mistral_api_key)
-embeddings  = OpenAIEmbeddings()
-############## VECTORSTORE ##################
-# vectorstore = FAISS.from_documents(
-#     documents=doc_splits,
-#     embedding=embeddings
+# docs = [WebBaseLoader(url).load() for url in urls]
+# docs_list = [item for sublist in docs for item in sublist]
+
+# text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+#     chunk_size=250, chunk_overlap=0
 # )
-vectorstore = Chroma.from_documents(
-    documents=doc_splits,
-    collection_name="rag-chroma",
-    embedding=embeddings
-)
+# doc_splits = text_splitter.split_documents(docs_list)
+
+####### PDF
+def load_chunk_persist_pdf() -> Chroma:
+    pdf_folder_path = "data/pdf_folder/"
+    documents = []
+    for file in os.listdir(pdf_folder_path):
+        if file.endswith('.pdf'):
+            pdf_path = os.path.join(pdf_folder_path, file)
+            loader = PyPDFLoader(pdf_path)
+            documents.extend(loader.load())
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=10)
+    chunked_documents = text_splitter.split_documents(documents)
+
+    vectorstore = Chroma.from_documents(
+        documents=chunked_documents,
+        embedding=MistralAIEmbeddings(),
+        persist_directory="data/chroma_store/"
+    )
+    vectorstore.persist()
+    return vectorstore
+
+# from langchain_community.document_loaders import PyPDFLoader
+# loader = PyPDFLoader('data/fitness_programs/ZeroToHero.pdf')
+# pages = loader.load_and_split()
+
+# text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+# splits = text_splitter.split_documents(pages)
+# vectorstore = Chroma.from_documents(documents=splits, embedding=MistralAIEmbeddings())
+vectorstore = load_chunk_persist_pdf()
 retriever = vectorstore.as_retriever()
+prompt = hub.pull("rlm/rag-prompt")
+
 
 # Data model
 class RouteQuery(BaseModel):
@@ -61,21 +74,35 @@ class RouteQuery(BaseModel):
     )
 
 # LLM with function call 
-# llm = ChatMistralAI(model="mistral-large-latest", mistral_api_key=mistral_api_key, temperature=0)
-
-# structured_llm_router = llm.with_structured_output(RouteQuery)
-
-# # Prompt 
-# system = """You are an expert at routing a user question to a vectorstore or web search.
-# The vectorstore contains documents related to agents, prompt engineering, and adversarial attacks.
-# Use the vectorstore for questions on these topics. For all else, use web-search."""
-# route_prompt = ChatPromptTemplate.from_messages(
-#     [
-#         ("system", system),
-#         ("human", "{question}"),
-#     ]
-# )
+llm = ChatMistralAI(model="mistral-large-latest", mistral_api_key=mistral_api_key, temperature=0)
+
+# structured_llm_router = llm.with_structured_output(RouteQuery, method="json_mode")
+
+# Prompt 
+system = """You are an expert at routing a user question to a vectorstore or web search.
+The vectorstore contains documents related to agents, prompt engineering, and adversarial attacks.
+Use the vectorstore for questions on these topics. For all else, use web-search."""
+route_prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", system),
+        ("human", "{question}"),
+    ]
+)
+prompt = hub.pull("rlm/rag-prompt")
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+
+def format_docs(docs):
+        return "\n\n".join(doc.page_content for doc in docs)
+    
+    
+rag_chain = (
+        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        | prompt
+        | llm
+        | StrOutputParser()
+    )
+
+print(rag_chain.invoke("Build a fitness program for me. Be precise in terms of exercises"))
 
-# question_router = route_prompt | structured_llm_router
-# print(question_router.invoke({"question": "Who will the Bears draft first in the NFL draft?"}))
-# print(question_router.invoke({"question": "What are the types of agent memory?"}))
+# print(rag_chain.invoke("I am a 45 years old woman and I have to loose weight for the summer. Provide me with a fitness program"))

Modules/websearch_agent.py

@@ -0,0 +1,30 @@
+import os
+os.environ['TOKENIZERS_PARALLELISM'] = 'true'
+os.environ['MISTRAL_API_KEY'] = "i5jSJkCFNGKfgIztloxTMjfckiFbYBj4"
+
+from langchain import hub
+from langchain.agents import AgentExecutor, create_json_chat_agent
+from langchain_mistralai.chat_models import ChatMistralAI
+
+prompt = hub.pull("hwchase17/react-chat-json")
+
+from langchain_community.tools import DuckDuckGoSearchRun
+
+tools = [DuckDuckGoSearchRun()]
+
+llm = ChatMistralAI(model='mistral-large-latest')
+
+agent = create_json_chat_agent(
+    llm=llm,
+    tools=tools,
+    prompt=prompt,
+    )
+
+agent_executor = AgentExecutor(
+    agent=agent,
+    tools=tools, 
+    verbose=True, 
+    handle_parsing_errors=True
+)
+
+agent_executor.invoke({"input":"How many proteins should I eat per day? Search mainly on wikipedia"})