Added rag dependencies

Author: Maximilian-Winter

ReadMe.md

@@ -1,7 +1,7 @@
 # llama-cpp-agent Framework
 
 ## Introduction
-The llama-cpp-agent framework is a tool designed for easy interaction with Large Language Models (LLMs). Allowing users to chat with LLM models, execute structured function calls and get structured output (objects).
+The llama-cpp-agent framework is a tool designed for easy interaction with Large Language Models (LLMs). Allowing users to chat with LLM models, execute structured function calls, get structured output (objects) and do retrieval augmented generation.
 
 It provides a simple yet robust interface and supports llama-cpp-python and OpenAI endpoints with GBNF grammar support (like the llama-cpp-python server) and the llama.cpp backend server.
 It works by generating a formal GGML-BNF grammar of the user defined structures and functions, which is then used by llama.cpp to generate text valid to that grammar. In contrast to most GBNF grammar generators it also supports nested objects, dictionaries, enums and lists of them.
@@ -10,6 +10,7 @@ It works by generating a formal GGML-BNF grammar of the user defined structures
 - **Simple Chat Interface**: Engage in seamless conversations with LLMs.
 - **Structured Output**: Get structured output (objects) from LLMs.
 - **Single and Parallel Function Calling**: Let the LLM execute functions.
+- **RAG - Retrieval Augmented Generation**: Perform retrieval augmented generation with colbert reranking.
 - **Flexibility**: Suited for various applications from casual chatting to specific function executions.
 
 ## Installation
@@ -75,7 +76,6 @@ while True:
 
 ```python
 # Example that uses the FunctionCallingAgent class to create a function calling agent.
-# Example that uses the FunctionCallingAgent class to create a function calling agent.
 import datetime
 from enum import Enum
 from typing import Union, Optional
@@ -111,7 +111,7 @@ class MathOperation(Enum):
 
 
 # llama-cpp-agent also supports "Instructor" library like function definitions as Pydantic models for function calling.
-# Simple pydantic calculator tool for the agent that can add, subtract, multiply, and divide. Docstring and description of fields will be used in system prompt.
+# Simple pydantic calculator tool for the agent that can add, subtract, multiply, and divide. Docstring and description of fields will be used in the system prompt.
 class calculator(BaseModel):
     """
     Perform a math operation on two numbers.
@@ -299,6 +299,117 @@ title='The Feynman Lectures on Physics' author='Richard Feynman, Robert B. Leigh
 
 ```
 
+### RAG - Retrieval Augmented Generation
+This example shows how to do RAG with colbert reranking.
+```python
+import json
+
+from ragatouille.utils import get_wikipedia_page
+
+from llama_cpp_agent.messages_formatter import MessagesFormatterType
+
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from llama_cpp_agent.llm_agent import LlamaCppAgent
+from llama_cpp_agent.gbnf_grammar_generator.gbnf_grammar_from_pydantic_models import (
+    generate_gbnf_grammar_and_documentation,
+)
+from llama_cpp_agent.providers.llama_cpp_endpoint_provider import (
+    LlamaCppEndpointSettings,
+)
+from llama_cpp_agent.rag.rag_colbert_reranker import RAGColbertReranker
+from llama_cpp_agent.rag.text_utils import RecursiveCharacterTextSplitter
+
+
+# Initialize the chromadb vector database with a colbert reranker.
+rag = RAGColbertReranker(persistent=False)
+
+# Initialize a recursive character text splitter with the correct chunk size of the embedding model.
+length_function = len
+splitter = RecursiveCharacterTextSplitter(
+    separators=["\n\n", "\n", " ", ""],
+    chunk_size=512,
+    chunk_overlap=0,
+    length_function=length_function,
+    keep_separator=True
+)
+
+# Use the ragatouille helper function to get the content of a wikipedia page.
+page = get_wikipedia_page("Synthetic_diamond")
+
+# Split the text of the wikipedia page into chunks for the vector database.
+splits = splitter.split_text(page)
+
+# Add the splits into the vector database
+for split in splits:
+    rag.add_document(split)
+
+# Define the query we want to ask based on the retrieved information
+query = "What is a BARS apparatus?"
+
+# Define a pydantic class to represent a query extension as additional queries to the original query.
+class QueryExtension(BaseModel):
+    """
+    Represents an extension of a query as additional queries.
+    """
+    queries: List[str] = Field(default_factory=list, description="List of queries.")
+
+
+# Generate a grammar and documentation of the query extension model.
+grammar, docs = generate_gbnf_grammar_and_documentation([QueryExtension])
+
+# Define a llamacpp server endpoint.
+main_model = LlamaCppEndpointSettings(completions_endpoint_url="http://127.0.0.1:8080/completion")
+
+# Define a query extension agent which will extend the query with additional queries.
+query_extension_agent = LlamaCppAgent(
+    main_model,
+    debug_output=True,
+    system_prompt="You are a world class query extension algorithm capable of extending queries by writing new queries. Do not answer the queries, simply provide a list of additional queries in JSON format. Structure your output according to the following model:\n\n" + docs.strip(),
+    predefined_messages_formatter_type=MessagesFormatterType.MIXTRAL
+)
+
+# Perform the query extension with the agent.
+output = query_extension_agent.get_chat_response(
+    f"Consider the following query: {query}", grammar=grammar)
+
+# Load the query extension in JSON format and create an instance of the query extension model.
+queries = QueryExtension.model_validate(json.loads(output))
+
+# Define the final prompt for the query with the retrieved information
+prompt = "Consider the following context:\n==========Context===========\n"
+
+# Retrieve the most fitting document chunks based on the original query and add them to the prompt.
+documents = rag.retrieve_documents(query, k=3)
+for doc in documents:
+    prompt += doc["content"] + "\n\n"
+
+# Retrieve the most fitting document chunks based on the extended queries and add them to the prompt.
+for qu in queries.queries:
+    documents = rag.retrieve_documents(qu, k=3)
+    for doc in documents:
+        if doc["content"] not in prompt:
+            prompt += doc["content"] + "\n\n"
+prompt += "\n======================\nQuestion: " + query
+
+# Define a new agent to answer the original query based on the retrieved information.
+agent_with_rag_information = LlamaCppAgent(
+    main_model,
+    debug_output=True,
+    system_prompt="You are an advanced AI assistant, trained by OpenAI. Only answer question based on the context information provided.",
+    predefined_messages_formatter_type=MessagesFormatterType.MIXTRAL
+)
+
+# Ask the agent the original query with the generated prompt that contains the retrieved information.
+agent_with_rag_information.get_chat_response(prompt)
+
+```
+Example output
+```text
+ BARS (Bridgman-Anvil High Pressure Reactor System) apparatus is a type of diamond-producing press used in the HPHT (High Pressure High Temperature) method for synthetic diamond growth. It consists of a ceramic cylindrical "synthesis capsule" placed in a cube of pressure-transmitting material, which is pressed by inner anvils and outer anvils. The whole assembly is locked in a disc-type barrel filled with oil, which pressurizes upon heating, and the oil pressure is transferred to the central cell. The BARS apparatus is claimed to be the most compact, efficient, and economical press design for diamond synthesis.
+```
 
 ### Knowledge Graph Creation Example
 This example, based on an example of the Instructor library for OpenAI,

examples/06_Rag/example_synthetic_diamonds_bars.py

@@ -75,7 +75,7 @@ class QueryExtension(BaseModel):
 query_extension_agent = LlamaCppAgent(
     main_model,
     debug_output=True,
-    system_prompt="You are a world class query extension algorithm capable of extending queries by writing new queries. Do not answer the queries, simply provide a list of additional queries in JSON format. Structure your output according to the following model:\n\n" + docs,
+    system_prompt="You are a world class query extension algorithm capable of extending queries by writing new queries. Do not answer the queries, simply provide a list of additional queries in JSON format. Structure your output according to the following model:\n\n" + docs.strip(),
     predefined_messages_formatter_type=MessagesFormatterType.MIXTRAL
 )
 
@@ -90,13 +90,13 @@ class QueryExtension(BaseModel):
 prompt = "Consider the following context:\n==========Context===========\n"
 
 # Retrieve the most fitting document chunks based on the original query and add them to the prompt.
-documents = rag.retrieve_documents(query, k=2)
+documents = rag.retrieve_documents(query, k=3)
 for doc in documents:
     prompt += doc["content"] + "\n\n"
 
 # Retrieve the most fitting document chunks based on the extended queries and add them to the prompt.
 for qu in queries.queries:
-    documents = rag.retrieve_documents(qu, k=2)
+    documents = rag.retrieve_documents(qu, k=3)
     for doc in documents:
         if doc["content"] not in prompt:
             prompt += doc["content"] + "\n\n"
@@ -106,7 +106,7 @@ class QueryExtension(BaseModel):
 agent_with_rag_information = LlamaCppAgent(
     main_model,
     debug_output=True,
-    system_prompt="You are an advanced AI assistant, trained by OpenAI.",
+    system_prompt="You are an advanced AI assistant, trained by OpenAI. Only answer question based on the context information provided.",
     predefined_messages_formatter_type=MessagesFormatterType.MIXTRAL
 )
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "llama-cpp-agent"
-version = "0.0.26"
+version = "0.0.27"
 description = "A framework for building LLM based AI agents with llama-cpp-python."
 
 readme = "ReadMe.md"
@@ -24,6 +24,7 @@ email = "[email protected]"
 
 [project.optional-dependencies]
 agent_memory = ["chromadb", "SQLAlchemy", "numpy", "scipy"]
+rag = ["ragatouille"]
 
 [project.urls]
 Homepage = "https://github.com/Maximilian-Winter/llama-cpp-agent"