├── utils ├── __init__.py ├── prompts.py ├── tools.py └── utils.py ├── data ├── shap.pdf └── file.txt ├── .gitignore ├── requirements.txt ├── readme.md └── agent.py /utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /data/shap.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/maxidiazbattan/Agentic-RAG-llamaindex-ollama/HEAD/data/shap.pdf -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | .lightning_studio 3 | .ollama 4 | .idea/this_studio.iml 5 | .vscode/settings.json 6 | .bash_history 7 | data/db -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | streamlit==1.31.0 2 | pydantic 3 | ollama 4 | llama-index 5 | llama-index-llms-ollama 6 | llama-index-embeddings-huggingface 7 | -------------------------------------------------------------------------------- /utils/prompts.py: -------------------------------------------------------------------------------- 1 | context = """ 2 | Purpose: You're an AI agent and your role is to provide accurate information about the documents provided. 3 | Please check your tools available before answer any question. 4 | """ -------------------------------------------------------------------------------- /utils/tools.py: -------------------------------------------------------------------------------- 1 | # tools 2 | import os 3 | 4 | 5 | # file saver tool 6 | file = os.path.join("data", "file.txt") 7 | 8 | def save_file(text): 9 | if not os.path.exists(file): 10 | with open(file, "w") as f: 11 | pass 12 | 13 | with open(file, "a") as f: 14 | f.writelines([text + "\n"]) 15 | 16 | return "file saved" 17 | 18 | -------------------------------------------------------------------------------- /data/file.txt: -------------------------------------------------------------------------------- 1 | The user requested a summary of a document consisting of research papers and articles focusing on explaining black box models in machine learning. The summary includes information about various methods for summarizing machine learning models and explainable AI (XAI) applications, as well as comparisons between different feature importance measures and their limitations. 2 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Agent with Llamaindex & Ollama 2 | 3 | This repo contains an implementation of an Agentic Retrieval-Augmented Generation (RAG) application leveraging Llamaindex and Ollama. The agent efficiently retrieves relevant information from documents and generates coherent, contextually appropriate responses, thanks to the tools provided. Additionally, you can save the responses to a file if desired. 4 | If you want, you can test this app in lightning studios via this [link](https://lightning.ai/maxidiazbattan/studios/agentic-rag-llamaindex-ollama), or locally following the steps below. 5 | 6 | ### 1. [Install](https://github.com/ollama/ollama?tab=readme-ov-file) ollama and pull models 7 | 8 | On linux 9 | ```shell 10 | curl -fsSL https://ollama.com/install.sh | sh 11 | ``` 12 | 13 | Start Ollama 14 | 15 | ```shell 16 | ollama serve 17 | ``` 18 | 19 | Pull the LLM you'd like to use: 20 | 21 | ```shell 22 | ollama pull koesn/mistral-7b-instruct 23 | ``` 24 | 25 | ### 2. Create a virtual environment 26 | 27 | ```shell 28 | python -m venv venv 29 | source venv/bin/activate 30 | ``` 31 | 32 | ### 3. Install libraries 33 | 34 | ```shell 35 | pip install -r requirements.txt 36 | ``` 37 | 38 | ### 4. Download the pdf to RAG, in this case a paper about Shap values 39 | ```shell 40 | wget https://arxiv.org/pdf/2302.08160 -O ./data/shap.pdf 41 | ``` 42 | 43 | ### 5. Run the agent 44 | 45 | ```shell 46 | python agent.py 47 | ``` 48 | -------------------------------------------------------------------------------- /agent.py: -------------------------------------------------------------------------------- 1 | # tools 2 | import os 3 | import glob 4 | 5 | # utils 6 | from utils.utils import DocumentToolsGenerator 7 | from utils.prompts import context 8 | 9 | # llamaindex 10 | from llama_index.embeddings.huggingface import HuggingFaceEmbedding 11 | from llama_index.llms.ollama import Ollama 12 | from llama_index.core import Settings 13 | from llama_index.core.agent import ReActAgent 14 | 15 | # warnings removal 16 | from warnings import filterwarnings 17 | filterwarnings('ignore') 18 | 19 | 20 | path = './data' 21 | file_name = 'shap.pdf' 22 | file_path = os.path.join(path, file_name) 23 | 24 | models = { 25 | 'mistral': 'koesn/mistral-7b-instruct', 26 | 'misal': 'smallstepai/misal-7B-instruct-v0.1', 27 | 'codellama': 'codellama' 28 | } 29 | 30 | # global settings 31 | embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5") 32 | llm = Ollama(model=models['mistral']) 33 | 34 | Settings.embed_model = embed_model 35 | Settings.llm = llm 36 | 37 | # DocumentToolsGenerator class instantiation 38 | docs_tools = DocumentToolsGenerator(file_path=file_path) 39 | 40 | # nodes creation 41 | nodes = docs_tools.data_ingestion() 42 | 43 | # tool generation 44 | vector_tool, summary_tool, file_tool = docs_tools.tool_generator(nodes=nodes) 45 | 46 | # agent initialization 47 | agent = ReActAgent.from_tools(tools=[vector_tool, summary_tool, file_tool], llm=llm, context=context, verbose=True) 48 | 49 | while (prompt := input("Enter a prompt (q to quit): ")) != "q": 50 | result = agent.query(prompt) 51 | print(result) 52 | -------------------------------------------------------------------------------- /utils/utils.py: -------------------------------------------------------------------------------- 1 | # tools 2 | import os 3 | from typing import Optional, List, Tuple 4 | 5 | # utils 6 | from utils.tools import save_file 7 | 8 | # llamaindex 9 | from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, StorageContext, SummaryIndex 10 | from llama_index.core.node_parser import SentenceSplitter 11 | from llama_index.core.schema import BaseNode 12 | from llama_index.core.tools import FunctionTool,QueryEngineTool 13 | from llama_index.core.vector_stores import MetadataFilters, FilterCondition 14 | 15 | 16 | class DocumentToolsGenerator: 17 | """ 18 | Document processing and tool generation for vector search and summarization. 19 | 20 | """ 21 | 22 | def __init__(self, file_path: str): 23 | """ 24 | Initializes the DocumentToolsGenerator with the given file path. 25 | 26 | Args: 27 | file_path (str): The path to the document file to be processed. 28 | """ 29 | self.file_path = file_path 30 | self.vector_index = None 31 | 32 | def data_ingestion(self, chunk_size: int = 1024, chunk_overlap: int = 64) -> List[BaseNode]: 33 | """ 34 | Loads and splits the document into chunks for processing. 35 | 36 | Args: 37 | chunk_size (int): The size of each chunk. Default is 1024. 38 | chunk_overlap (int): The overlap between chunks. Default is 64. 39 | 40 | Returns: 41 | List[BaseNode]: A list of document nodes created from the chunks. 42 | """ 43 | documents = SimpleDirectoryReader(input_files=[self.file_path]).load_data() 44 | sentence_splitter = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) 45 | nodes = sentence_splitter.get_nodes_from_documents(documents=documents) 46 | 47 | return nodes 48 | 49 | def vector_query(self, query: str, page_numbers: Optional[List[str]] = None) -> str: 50 | """ 51 | Performs a vector search over the document using the specified query and optional page numbers. 52 | 53 | Args: 54 | query (str): The query string to be embedded for the search. 55 | page_numbers (Optional[List[str]]): A list of page numbers to be retrieved. Defaults to searching over all pages. 56 | 57 | Returns: 58 | str: The search response. 59 | """ 60 | page_numbers = page_numbers or [] 61 | metadata_dict = [{"key": 'page_label', "value": p} for p in page_numbers] 62 | query_engine = self.vector_index.as_query_engine( 63 | similarity_top_k=2, 64 | filters=MetadataFilters.from_dicts( 65 | metadata_dict, 66 | condition=FilterCondition.OR) 67 | ) 68 | response = query_engine.query(query) 69 | return response 70 | 71 | def tool_generator(self, nodes: List[BaseNode], vector_store_path: str = './data/db', db_name: str = 'vs') -> Tuple[FunctionTool, QueryEngineTool, FunctionTool]: 72 | """ 73 | Generates and returns tools for vector search, document summarization, and file saving. 74 | 75 | Args: 76 | nodes (List[BaseNode]): The list of nodes generated from the document. 77 | vector_store_path (str): The path to store the vector index. Default is './data/db'. 78 | db_name (str): The name of the vector store database. Default is 'vs'. 79 | 80 | Returns: 81 | Tuple[FunctionTool, QueryEngineTool, FunctionTool]: A tuple containing the vector query tool, summary query tool, and file saving tool. 82 | """ 83 | # vector index 84 | if not os.path.exists(db_name): 85 | self.vector_index = VectorStoreIndex(nodes=nodes) 86 | self.vector_index.storage_context.vector_store.persist(persist_path=f'{vector_store_path}/{db_name}') 87 | else: 88 | self.vector_index = load_index_from_storage( 89 | StorageContext.from_defaults(persist_dir=f'{vector_store_path}') 90 | ) 91 | 92 | # summary index 93 | summary_index = SummaryIndex(nodes=nodes) 94 | 95 | # prepare vector tool 96 | vector_query_tool = FunctionTool.from_defaults( 97 | name="vector_search_tool", 98 | fn=self.vector_query, 99 | description="Useful for searching specific facts in a document" 100 | ) 101 | 102 | # prepare summary tool 103 | summary_query = summary_index.as_query_engine(response_mode="tree_summarize") 104 | summary_query_tool = QueryEngineTool.from_defaults( 105 | name="summary_query_tool", 106 | query_engine=summary_query, 107 | description="Useful for summarizing an entire document. DO NOT USE if you have specified questions over the documents." 108 | ) 109 | 110 | # prepare file saving tool 111 | file_tool = FunctionTool.from_defaults( 112 | name="file_saver_tool", 113 | fn=save_file, 114 | description="Useful for saving a text file" 115 | ) 116 | 117 | return vector_query_tool, summary_query_tool, file_tool 118 | --------------------------------------------------------------------------------