OPEN AI - RAG
Simon-Pierre Boucher
2024-09-14
This Python script integrates OpenAI's API to build a chatbot capable of retrieving relevant information from a set of documents and using that context to generate a more informed response. Here's a breakdown of its functionality:
1. Environment Setup:¶
load_dotenv(): Loads environment variables, including the OpenAI API key, from a.envfile.- The API key is fetched from environment variables using
os.getenv("OPENAI_API_KEY").
2. retrieve_relevant_docs() Function:¶
- This function scans a list of documents to retrieve those containing keywords relevant to the current user query.
- It checks for keywords in the user query (splitting it by spaces) and returns documents that match any of those keywords.
3. make_openai_api_call_with_rag() Function:¶
- This function combines the retrieved relevant documents with the conversation history and the user’s current message to create a more contextually aware prompt.
- The steps:
- It retrieves relevant documents using
retrieve_relevant_docs(). - It constructs the context from the retrieved documents.
- It augments the user's current message with the relevant context.
- It sends the entire conversation history, including the augmented message, to the OpenAI API via the
/chat/completionsendpoint.
- It retrieves relevant documents using
- The function sends the API request with customizable parameters, such as
model,temperature,max_tokens, etc. - The response is returned as a JSON object.
4. format_markdown() Function:¶
- Converts the API response, which might contain Markdown, into HTML for better display in a Jupyter notebook.
- It handles common Markdown features like bold (
**text**), italic (*text*), and headings (e.g.,# Heading,## Sub-heading).
5. format_response() Function:¶
- This function formats the OpenAI API's JSON response for HTML display.
- It extracts key details such as:
- The model used for generating the response.
- Token usage statistics (prompt, completion, and total tokens).
- The assistant's response in a well-formatted HTML block, including the role and content of the message.
- Finally, it uses
IPython.display.display()to render the response as an HTML block inside the notebook.
6. Example Usage:¶
- The script starts with a basic conversation history that includes an assistant greeting.
- The current user query is: "What is the debt-to-equity ratio of ABC Corp?"
- The function searches through a predefined list of documents containing information about ABC Corp.
- Relevant documents (those mentioning the debt-to-equity ratio) are retrieved, and this context is added to the user query.
- The combined history and query are sent to OpenAI's API to generate a more informed and specific response.
- The response is then formatted and displayed in the notebook.
Example Workflow:¶
- The user query is: "What is the debt-to-equity ratio of ABC Corp?"
- The script searches through the documents and finds: "ABC Corp. has a current debt-to-equity ratio of 0.3."
- The query, along with this relevant document, is sent to OpenAI's API.
- The API generates a response based on the conversation history and relevant documents.
- The formatted response is displayed in HTML.
Result:¶
The script augments OpenAI's basic conversational abilities with a simple RAG (Retrieval-Augmented Generation) system, making it capable of answering user queries with context from specific, relevant documents. This is especially useful in tasks like document analysis, question answering, and information retrieval in a structured environment.
In [1]:
import os
import requests
from dotenv import load_dotenv
from IPython.display import display, HTML
import re
# Load environment variables from the .env file
load_dotenv()
# Get the API key from environment variables
api_key = os.getenv("OPENAI_API_KEY")
def retrieve_relevant_docs(query, documents):
"""
Simple keyword-based function to retrieve relevant documents.
"""
relevant_docs = []
for doc in documents:
if any(keyword.lower() in doc.lower() for keyword in query.split()):
relevant_docs.append(doc)
return relevant_docs
def make_openai_api_call_with_rag(conversation_history, current_message, documents, model="gpt-4o", temperature=1, max_tokens=256, top_p=1, frequency_penalty=0, presence_penalty=0):
"""
Makes an OpenAI API call using the provided conversation history and current message
with retrieved relevant documents.
:param conversation_history: List of conversation history messages
:param current_message: Current user message
:param documents: List of available documents for retrieval
:param model: OpenAI model to use
:param temperature: Controls the creativity of the response
:param max_tokens: Maximum number of tokens in the response
:param top_p: Probability filtering
:param frequency_penalty: Frequency penalty
:param presence_penalty: Presence penalty
:return: JSON response from the OpenAI API
"""
# Retrieve relevant documents based on the current message
relevant_docs = retrieve_relevant_docs(current_message, documents)
# Combine the retrieved documents into a single context
context = "\n\n".join(relevant_docs)
# Add the retrieved context to the current message
augmented_message = f"Context: {context}\n\n{current_message}"
# Combine the history and the augmented current message
messages = conversation_history + [{"role": "user", "content": augmented_message}]
url = 'https://api.openai.com/v1/chat/completions'
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {api_key}'
}
data = {
"model": model,
"messages": messages,
"temperature": temperature,
"max_tokens": max_tokens,
"top_p": top_p,
"frequency_penalty": frequency_penalty,
"presence_penalty": presence_penalty
}
response = requests.post(url, headers=headers, json=data)
return response.json()
def format_markdown(content):
"""
Converts Markdown content to HTML.
:param content: Markdown text
:return: HTML text
"""
# Remove unnecessary line breaks after enumerations
content = re.sub(r'(\d+\..*?)\n\n', r'\1\n', content)
# Convert Markdown to HTML
content = content.replace('\n', '<br>')
content = re.sub(r'\*\*(.*?)\*\*', r'<strong>\1</strong>', content)
content = re.sub(r'\*(.*?)\*', r'<em>\1</em>', content)
content = re.sub(r'### (.*)', r'<h3>\1</h3>', content)
content = re.sub(r'## (.*)', r'<h2>\1</h2>', content)
content = re.sub(r'# (.*)', r'<h1>\1</h1>', content)
return content
def format_response(response):
"""
Formats the JSON response from the OpenAI API for HTML display.
:param response: JSON response from the OpenAI API
:return: None
"""
result = response['choices'][0]['message']
usage = response['usage']
# Format content with Markdown
formatted_content = format_markdown(result['content'])
html = """
<div class="api-response">
"""
# Model Information
html += f"""
<div class="bubble">
<h3>Model Information</h3>
<p><strong>Model:</strong> {response['model']}</p>
</div>
"""
# Token Usage
html += f"""
<div class="bubble">
<h3>Token Usage</h3>
<p><strong>Prompt Tokens:</strong> {usage['prompt_tokens']}</p>
<p><strong>Completion Tokens:</strong> {usage['completion_tokens']}</p>
<p><strong>Total Tokens:</strong> {usage['total_tokens']}</p>
</div>
"""
# Response Content
html += f"""
<div class="bubble">
<h3>Response Content</h3>
<p><strong>Role:</strong> {result['role']}</p>
<p><strong>Content:</strong></p>
<div>{formatted_content}</div>
</div>
"""
html += "</div>"
display(HTML(html))
In [2]:
conversation_history = [
{"role": "assistant", "content": "Hello! How can I assist you today?"}
]
current_message = "What is the debt-to-equity ratio of ABC Corp?"
documents = [
"ABC Corp. reported a revenue of 50 million for Q2 2024, a 10 percent increase from Q1 2024. The company's net income for the quarter was 5 million, reflecting a 5 percent profit margin.",
"ABC Corp. has a current debt-to-equity ratio of 0.3, indicating that the company has a low level of debt compared to its equity.",
"The market capitalization of ABC Corp. is currently 300 million, based on a share price of 30 and 10 million shares outstanding.",
"In Q2 2024, ABC Corp. announced a dividend of 0.50 per share, which will be distributed to shareholders on October 1, 2024.",
"ABC Corp.'s gross profit margin for Q2 2024 was 40 percent, reflecting strong control over cost of goods sold and efficient operations."
]
In [3]:
response = make_openai_api_call_with_rag(conversation_history, current_message, documents,model="gpt-4o")
format_response(response)
In [4]:
response = make_openai_api_call_with_rag(conversation_history, current_message, documents,model="gpt-4o-mini")
format_response(response)
In [5]:
response = make_openai_api_call_with_rag(conversation_history, current_message, documents,model="gpt-4")
format_response(response)
In [6]:
response = make_openai_api_call_with_rag(conversation_history, current_message, documents,model="gpt-4-turbo")
format_response(response)
In [7]:
response = make_openai_api_call_with_rag(conversation_history, current_message, documents,model="gpt-3.5-turbo")
format_response(response)