forked from cleveragents/cleveragents-core
aditya
1b8fba684f
User-triggered simulation via terminal commands. RAG app with OpenAI embeddings, FAISS vector store, document chunking, and CLI interface.
376 lines
12 KiB
Python
376 lines
12 KiB
Python
import os
|
|
import argparse
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import List, Tuple, Optional
|
|
import numpy as np
|
|
import faiss
|
|
import tiktoken
|
|
import openai
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RAGApplication:
|
|
"""
|
|
A complete RAG (Retrieval-Augmented Generation) application that processes documents,
|
|
creates embeddings, stores them in FAISS, and answers questions using OpenAI GPT.
|
|
"""
|
|
|
|
def __init__(self, model_name: str = "gpt-3.5-turbo", embedding_model: str = "text-embedding-ada-002"):
|
|
"""
|
|
Initialize the RAG application.
|
|
|
|
Args:
|
|
model_name: OpenAI model name for chat completions
|
|
embedding_model: OpenAI model name for embeddings
|
|
"""
|
|
self.model_name = model_name
|
|
self.embedding_model = embedding_model
|
|
self.chunk_size = 1000
|
|
self.chunk_overlap = 200
|
|
self.max_chunks_per_query = 5
|
|
|
|
# Initialize OpenAI client
|
|
api_key = os.getenv("OPENAI_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("OPENAI_API_KEY environment variable not found")
|
|
|
|
self.client = openai.OpenAI(api_key=api_key)
|
|
|
|
# Initialize tokenizer
|
|
try:
|
|
self.tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
|
|
except KeyError:
|
|
self.tokenizer = tiktoken.get_encoding("cl100k_base")
|
|
|
|
# Storage for documents and embeddings
|
|
self.documents: List[str] = []
|
|
self.chunks: List[str] = []
|
|
self.embeddings: Optional[np.ndarray] = None
|
|
self.faiss_index: Optional[faiss.Index] = None
|
|
|
|
def load_documents(self, docs_dir: str) -> None:
|
|
"""
|
|
Load all .txt files from the specified directory.
|
|
|
|
Args:
|
|
docs_dir: Path to directory containing .txt files
|
|
"""
|
|
docs_path = Path(docs_dir)
|
|
|
|
if not docs_path.exists():
|
|
raise FileNotFoundError(f"Directory not found: {docs_dir}")
|
|
|
|
if not docs_path.is_dir():
|
|
raise ValueError(f"Path is not a directory: {docs_dir}")
|
|
|
|
txt_files = list(docs_path.glob("*.txt"))
|
|
|
|
if not txt_files:
|
|
raise ValueError(f"No .txt files found in directory: {docs_dir}")
|
|
|
|
logger.info(f"Found {len(txt_files)} .txt files in {docs_dir}")
|
|
|
|
for txt_file in txt_files:
|
|
try:
|
|
with open(txt_file, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
if content.strip(): # Only add non-empty files
|
|
self.documents.append(content)
|
|
logger.info(f"Loaded: {txt_file.name}")
|
|
except Exception as e:
|
|
logger.error(f"Error loading {txt_file}: {e}")
|
|
continue
|
|
|
|
if not self.documents:
|
|
raise ValueError("No valid documents were loaded")
|
|
|
|
logger.info(f"Successfully loaded {len(self.documents)} documents")
|
|
|
|
def chunk_text(self, text: str) -> List[str]:
|
|
"""
|
|
Split text into chunks with overlap for better context preservation.
|
|
|
|
Args:
|
|
text: Text to chunk
|
|
|
|
Returns:
|
|
List of text chunks
|
|
"""
|
|
tokens = self.tokenizer.encode(text)
|
|
chunks = []
|
|
|
|
for i in range(0, len(tokens), self.chunk_size - self.chunk_overlap):
|
|
chunk_tokens = tokens[i:i + self.chunk_size]
|
|
chunk_text = self.tokenizer.decode(chunk_tokens)
|
|
chunks.append(chunk_text)
|
|
|
|
return chunks
|
|
|
|
def process_documents(self) -> None:
|
|
"""
|
|
Process all loaded documents by chunking them into smaller pieces.
|
|
"""
|
|
if not self.documents:
|
|
raise ValueError("No documents loaded. Call load_documents() first.")
|
|
|
|
logger.info("Processing documents into chunks...")
|
|
|
|
self.chunks = []
|
|
for i, doc in enumerate(self.documents):
|
|
doc_chunks = self.chunk_text(doc)
|
|
self.chunks.extend(doc_chunks)
|
|
logger.info(f"Document {i+1}: created {len(doc_chunks)} chunks")
|
|
|
|
logger.info(f"Total chunks created: {len(self.chunks)}")
|
|
|
|
def get_embeddings(self, texts: List[str], batch_size: int = 100) -> np.ndarray:
|
|
"""
|
|
Get embeddings for a list of texts using OpenAI API.
|
|
|
|
Args:
|
|
texts: List of texts to embed
|
|
batch_size: Number of texts to process in each API call
|
|
|
|
Returns:
|
|
NumPy array of embeddings
|
|
"""
|
|
embeddings = []
|
|
|
|
for i in range(0, len(texts), batch_size):
|
|
batch = texts[i:i + batch_size]
|
|
|
|
try:
|
|
response = self.client.embeddings.create(
|
|
model=self.embedding_model,
|
|
input=batch
|
|
)
|
|
|
|
batch_embeddings = [data.embedding for data in response.data]
|
|
embeddings.extend(batch_embeddings)
|
|
|
|
logger.info(f"Generated embeddings for batch {i//batch_size + 1}/{(len(texts)-1)//batch_size + 1}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating embeddings for batch {i//batch_size + 1}: {e}")
|
|
raise
|
|
|
|
return np.array(embeddings, dtype=np.float32)
|
|
|
|
def create_vector_store(self) -> None:
|
|
"""
|
|
Create embeddings for all chunks and initialize FAISS index.
|
|
"""
|
|
if not self.chunks:
|
|
raise ValueError("No chunks available. Call process_documents() first.")
|
|
|
|
logger.info("Generating embeddings for chunks...")
|
|
self.embeddings = self.get_embeddings(self.chunks)
|
|
|
|
# Initialize FAISS index
|
|
embedding_dim = self.embeddings.shape[1]
|
|
self.faiss_index = faiss.IndexFlatIP(embedding_dim) # Inner product for cosine similarity
|
|
|
|
# Normalize embeddings for cosine similarity
|
|
faiss.normalize_L2(self.embeddings)
|
|
|
|
# Add embeddings to index
|
|
self.faiss_index.add(self.embeddings)
|
|
|
|
logger.info(f"Created FAISS index with {self.faiss_index.ntotal} vectors")
|
|
|
|
def retrieve_relevant_chunks(self, query: str, k: int = None) -> List[Tuple[str, float]]:
|
|
"""
|
|
Retrieve the most relevant chunks for a given query.
|
|
|
|
Args:
|
|
query: User query
|
|
k: Number of chunks to retrieve (default: self.max_chunks_per_query)
|
|
|
|
Returns:
|
|
List of tuples (chunk_text, similarity_score)
|
|
"""
|
|
if self.faiss_index is None:
|
|
raise ValueError("Vector store not initialized. Call create_vector_store() first.")
|
|
|
|
if k is None:
|
|
k = min(self.max_chunks_per_query, len(self.chunks))
|
|
|
|
# Get query embedding
|
|
query_embedding = self.get_embeddings([query])
|
|
faiss.normalize_L2(query_embedding)
|
|
|
|
# Search for similar chunks
|
|
similarities, indices = self.faiss_index.search(query_embedding, k)
|
|
|
|
# Return chunks with their similarity scores
|
|
results = []
|
|
for i, (similarity, idx) in enumerate(zip(similarities[0], indices[0])):
|
|
if idx < len(self.chunks): # Ensure valid index
|
|
results.append((self.chunks[idx], float(similarity)))
|
|
|
|
return results
|
|
|
|
def generate_answer(self, query: str, context_chunks: List[str]) -> str:
|
|
"""
|
|
Generate an answer using OpenAI GPT with the retrieved context.
|
|
|
|
Args:
|
|
query: User query
|
|
context_chunks: List of relevant text chunks
|
|
|
|
Returns:
|
|
Generated answer
|
|
"""
|
|
# Prepare context
|
|
context = "\n\n".join(context_chunks)
|
|
|
|
# Create prompt
|
|
prompt = f"""Based on the following context, please answer the question. If the answer cannot be found in the context, please say so.
|
|
|
|
Context:
|
|
{context}
|
|
|
|
Question: {query}
|
|
|
|
Answer:"""
|
|
|
|
try:
|
|
response = self.client.chat.completions.create(
|
|
model=self.model_name,
|
|
messages=[
|
|
{"role": "system", "content": "You are a helpful assistant that answers questions based on provided context."},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
max_tokens=1000,
|
|
temperature=0.1
|
|
)
|
|
|
|
return response.choices[0].message.content.strip()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating answer: {e}")
|
|
return f"Sorry, I encountered an error while generating the answer: {e}"
|
|
|
|
def answer_question(self, question: str) -> str:
|
|
"""
|
|
Complete pipeline to answer a question using RAG.
|
|
|
|
Args:
|
|
question: User question
|
|
|
|
Returns:
|
|
Generated answer with context information
|
|
"""
|
|
logger.info(f"Processing question: {question}")
|
|
|
|
# Retrieve relevant chunks
|
|
relevant_chunks = self.retrieve_relevant_chunks(question)
|
|
|
|
if not relevant_chunks:
|
|
return "Sorry, I couldn't find any relevant information to answer your question."
|
|
|
|
# Extract chunk texts and log retrieval info
|
|
chunk_texts = [chunk for chunk, score in relevant_chunks]
|
|
logger.info(f"Retrieved {len(chunk_texts)} relevant chunks")
|
|
|
|
# Generate answer
|
|
answer = self.generate_answer(question, chunk_texts)
|
|
|
|
# Add metadata about sources
|
|
source_info = f"\n\n[Answer based on {len(chunk_texts)} relevant text segments from the document collection]"
|
|
|
|
return answer + source_info
|
|
|
|
def initialize(self, docs_dir: str) -> None:
|
|
"""
|
|
Complete initialization of the RAG system.
|
|
|
|
Args:
|
|
docs_dir: Directory containing .txt files
|
|
"""
|
|
logger.info("Initializing RAG system...")
|
|
|
|
self.load_documents(docs_dir)
|
|
self.process_documents()
|
|
self.create_vector_store()
|
|
|
|
logger.info("RAG system initialized successfully!")
|
|
|
|
|
|
def main():
|
|
"""
|
|
Command-line interface for the RAG application.
|
|
"""
|
|
parser = argparse.ArgumentParser(description="RAG Question Answering System")
|
|
parser.add_argument(
|
|
"--docs-dir",
|
|
required=True,
|
|
help="Directory containing .txt documents to index"
|
|
)
|
|
parser.add_argument(
|
|
"--question",
|
|
help="Question to ask (if not provided, will enter interactive mode)"
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
default="gpt-3.5-turbo",
|
|
help="OpenAI model name for answer generation"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
# Initialize RAG application
|
|
rag = RAGApplication(model_name=args.model)
|
|
rag.initialize(args.docs_dir)
|
|
|
|
if args.question:
|
|
# Single question mode
|
|
answer = rag.answer_question(args.question)
|
|
print("\nQuestion:", args.question)
|
|
print("\nAnswer:", answer)
|
|
else:
|
|
# Interactive mode
|
|
print("\nRAG system ready! Type 'quit' or 'exit' to stop.")
|
|
print("=" * 50)
|
|
|
|
while True:
|
|
try:
|
|
question = input("\nEnter your question: ").strip()
|
|
|
|
if question.lower() in ['quit', 'exit', 'q']:
|
|
print("Goodbye!")
|
|
break
|
|
|
|
if not question:
|
|
print("Please enter a question.")
|
|
continue
|
|
|
|
answer = rag.answer_question(question)
|
|
print(f"\nAnswer: {answer}")
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n\nGoodbye!")
|
|
break
|
|
except EOFError:
|
|
print("\n\nGoodbye!")
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error(f"Application error: {e}")
|
|
print(f"Error: {e}")
|
|
return 1
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main()) |