File size: 5,504 Bytes
3c18172 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
"""
RAG Integration Example
This file demonstrates how to integrate the RAG system with your existing app.py
without modifying the original code. This shows a separate example that you can
use as a reference when you're ready to integrate RAG functionality.
Usage:
1. Run train_rag.py to build your knowledge base
2. Use this integration example as a template for modifying your app.py
"""
from dotenv import load_dotenv
from openai import OpenAI
import json
import os
import requests
import gradio as gr
from rag_utils import RAGSystem
import logging
# Set up logging
logging.basicConfig(level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
load_dotenv(override=True)
# Initialize the RAG system with error handling
rag = RAGSystem()
# Try to load the pre-trained index
try:
rag_index_path = "me/rag_index"
if os.path.exists(rag_index_path):
logger.info(f"Loading RAG index from {rag_index_path}")
rag.load_index(rag_index_path)
logger.info(f"Loaded {len(rag.chunks)} chunks from index")
else:
logger.warning(f"RAG index not found at {rag_index_path}. Run train_rag.py first.")
# Create a flag to indicate missing index
rag.index_loaded = False
except Exception as e:
logger.error(f"Error loading RAG index: {e}")
# Create a flag to indicate error
rag.index_loaded = False
else:
# Mark index as successfully loaded
rag.index_loaded = True
class MeWithRAG:
"""Example class showing how to integrate RAG with your existing Me class"""
def __init__(self):
self.openai = OpenAI(api_key=os.getenv("GOOGLE_API_KEY"),
base_url="https://generativelanguage.googleapis.com/v1beta/openai/")
self.name = "Sagarnil Das"
# No need to load the entire PDF and summary here!
# Just use a brief intro for the system prompt
self.intro = "I'm a software engineer and data scientist with expertise in AI and machine learning."
def system_prompt(self, query=None):
"""Dynamic system prompt that includes RAG context if a query is provided"""
system_prompt = f"You are acting as {self.name}. You are answering questions on {self.name}'s website, \
particularly questions related to {self.name}'s career, background, skills and experience. \
Your responsibility is to represent {self.name} for interactions on the website as faithfully as possible."
# If we have a query and RAG index is loaded, get relevant context
if query and getattr(rag, 'index_loaded', False):
try:
context = rag.get_context_for_query(query, top_k=3)
system_prompt += f"\n\n## Relevant Background Information:\n{context}\n\n"
except Exception as e:
logger.error(f"Error retrieving context: {e}")
system_prompt += f"\n\n## Brief Introduction:\n{self.intro}\n\n"
else:
# Just use a brief intro for the initial prompt
system_prompt += f"\n\n## Brief Introduction:\n{self.intro}\n\n"
system_prompt += f"With this context, please chat with the user, always staying in character as {self.name}."
return system_prompt
def chat(self, message, history):
"""Chat function that uses RAG to retrieve relevant context"""
try:
# First message includes base system prompt
messages = [{"role": "system", "content": self.system_prompt()}]
# Add conversation history
if isinstance(history, list) and all(isinstance(h, dict) for h in history):
messages.extend(history)
else:
for user_msg, assistant_msg in history:
messages.append({"role": "user", "content": user_msg})
messages.append({"role": "assistant", "content": assistant_msg})
# Add the user's message
messages.append({"role": "user", "content": message})
# Now get a dynamic system prompt with RAG context for this specific query
# Replace the initial system message with one that includes relevant context
try:
messages[0] = {"role": "system", "content": self.system_prompt(query=message)}
except Exception as e:
logger.error(f"Error generating system prompt: {e}")
# Keep original system prompt if there's an error
# Call the LLM
response = self.openai.chat.completions.create(
model="gemini-2.0-flash",
messages=messages
)
return response.choices[0].message.content
except Exception as e:
logger.error(f"Error in chat method: {e}")
return f"I apologize, but I encountered an error while processing your request. Please try again later."
# Example of how to use this
if __name__ == "__main__":
# This is just a demonstration - not meant to be run directly
me_rag = MeWithRAG()
gr.ChatInterface(me_rag.chat, type="messages").launch()
# When integrating with your actual app.py, you could:
# 1. Update the Me class to use RAG-based context retrieval
# 2. Keep all your existing tool-calling functionality
# 3. Use the RAG system to provide context-aware responses |