forked from cleveragents/cleveragents-core
aditya
ff52ef4f69
First simulation testing subdirectory output support. RAG app with separate modules: document_loader, chunker, embeddings, vector_store, retriever, generator — plus config.py and main.py CLI entry point. All Python files pass syntax check. 9/10 DoD checks pass. Also adds sim_test_commands.md with 9 new simulation definitions (SIM9-SIM17) for manual terminal testing.
228 lines
8.1 KiB
Python
228 lines
8.1 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from dotenv import load_dotenv
|
|
|
|
from config import RAGConfig, ModelConfig
|
|
from src.document_loader import DocumentLoader
|
|
from src.chunker import TextChunker
|
|
from src.embeddings import EmbeddingGenerator
|
|
from src.vector_store import FAISSVectorStore
|
|
from src.retriever import DocumentRetriever
|
|
from src.generator import AnswerGenerator
|
|
|
|
|
|
class RAGSystem:
|
|
"""Main RAG system orchestrator."""
|
|
|
|
def __init__(self, config: RAGConfig):
|
|
self.config = config
|
|
self.config.ensure_directories()
|
|
|
|
# Initialize components
|
|
self.embedding_generator = EmbeddingGenerator(
|
|
api_key=config.openai_api_key,
|
|
model=config.embedding_model,
|
|
batch_size=config.embedding_batch_size
|
|
)
|
|
|
|
# Set embedding dimension based on model
|
|
embedding_dim = ModelConfig.get_embedding_dimension(config.embedding_model)
|
|
self.vector_store = FAISSVectorStore(dimension=embedding_dim)
|
|
|
|
self.retriever = DocumentRetriever(self.embedding_generator, self.vector_store)
|
|
|
|
self.generator = AnswerGenerator(
|
|
api_key=config.openai_api_key,
|
|
model=config.chat_model,
|
|
temperature=config.temperature
|
|
)
|
|
|
|
def build_index(self, force_rebuild: bool = False):
|
|
"""Build or load the vector index."""
|
|
index_path = self.config.get_index_path()
|
|
metadata_path = self.config.get_metadata_path()
|
|
|
|
# Check if index exists and we don't need to rebuild
|
|
if not force_rebuild and self.vector_store.exists(index_path, metadata_path):
|
|
print("Loading existing index...")
|
|
try:
|
|
self.vector_store.load(index_path, metadata_path)
|
|
print(f"Loaded index with {len(self.vector_store.texts)} chunks")
|
|
return
|
|
except Exception as e:
|
|
print(f"Error loading index: {e}")
|
|
print("Rebuilding index...")
|
|
|
|
# Load documents
|
|
print(f"Loading documents from {self.config.documents_dir}...")
|
|
loader = DocumentLoader(self.config.documents_dir)
|
|
try:
|
|
documents = loader.load_documents()
|
|
except FileNotFoundError as e:
|
|
print(f"Error: {e}")
|
|
return
|
|
|
|
if not documents:
|
|
print("No .txt files found in documents directory")
|
|
return
|
|
|
|
print(f"Loaded {len(documents)} documents")
|
|
|
|
# Chunk documents
|
|
print("Chunking documents...")
|
|
chunker = TextChunker(
|
|
chunk_size=self.config.chunk_size,
|
|
overlap=self.config.chunk_overlap,
|
|
model=self.config.embedding_model
|
|
)
|
|
chunks = chunker.chunk_documents(documents)
|
|
print(f"Created {len(chunks)} chunks")
|
|
|
|
if not chunks:
|
|
print("No chunks created")
|
|
return
|
|
|
|
# Generate embeddings
|
|
print("Generating embeddings...")
|
|
texts = [chunk[0] for chunk in chunks]
|
|
sources = [chunk[1] for chunk in chunks]
|
|
|
|
try:
|
|
embeddings = self.embedding_generator.generate_embeddings(texts)
|
|
except Exception as e:
|
|
print(f"Error generating embeddings: {e}")
|
|
return
|
|
|
|
# Create index
|
|
print("Creating vector index...")
|
|
self.vector_store.create_index(embeddings, texts, sources)
|
|
|
|
# Save index
|
|
print("Saving index...")
|
|
try:
|
|
self.vector_store.save(index_path, metadata_path)
|
|
print("Index saved successfully")
|
|
except Exception as e:
|
|
print(f"Error saving index: {e}")
|
|
|
|
def query(self, question: str) -> str:
|
|
"""Process a single query."""
|
|
if not question.strip():
|
|
return "Please provide a valid question."
|
|
|
|
# Retrieve relevant chunks
|
|
try:
|
|
chunks = self.retriever.retrieve_with_threshold(
|
|
question,
|
|
top_k=self.config.top_k,
|
|
threshold=self.config.similarity_threshold
|
|
)
|
|
except Exception as e:
|
|
return f"Error during retrieval: {e}"
|
|
|
|
if not chunks:
|
|
return "I couldn't find relevant information to answer your question."
|
|
|
|
# Generate answer
|
|
try:
|
|
answer = self.generator.generate_answer(question, chunks)
|
|
return answer
|
|
except Exception as e:
|
|
return f"Error generating answer: {e}"
|
|
|
|
def interactive_mode(self):
|
|
"""Run interactive query mode."""
|
|
print("\n=== Interactive RAG System ===")
|
|
print("Ask questions about your documents. Type 'quit' or 'exit' to stop.\n")
|
|
|
|
while True:
|
|
try:
|
|
question = input("Question: ").strip()
|
|
|
|
if question.lower() in ['quit', 'exit', 'q']:
|
|
print("Goodbye!")
|
|
break
|
|
|
|
if not question:
|
|
continue
|
|
|
|
print("\nThinking...")
|
|
answer = self.query(question)
|
|
print(f"\nAnswer: {answer}\n")
|
|
print("-" * 50)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nGoodbye!")
|
|
break
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
|
|
|
|
def main():
|
|
"""Main CLI entry point."""
|
|
load_dotenv()
|
|
|
|
parser = argparse.ArgumentParser(description="RAG (Retrieval-Augmented Generation) System")
|
|
parser.add_argument("--docs-dir", default="documents", help="Directory containing .txt documents")
|
|
parser.add_argument("--index-dir", default="index", help="Directory for storing vector index")
|
|
parser.add_argument("--question", "-q", help="Single question to ask")
|
|
parser.add_argument("--interactive", "-i", action="store_true", help="Run in interactive mode")
|
|
parser.add_argument("--rebuild-index", action="store_true", help="Force rebuild of vector index")
|
|
parser.add_argument("--embedding-model", default="text-embedding-3-small",
|
|
choices=list(ModelConfig.EMBEDDING_MODELS.keys()),
|
|
help="OpenAI embedding model to use")
|
|
parser.add_argument("--chat-model", default="gpt-3.5-turbo",
|
|
choices=ModelConfig.CHAT_MODELS,
|
|
help="OpenAI chat model to use")
|
|
parser.add_argument("--top-k", type=int, default=5, help="Number of chunks to retrieve")
|
|
parser.add_argument("--chunk-size", type=int, default=1000, help="Chunk size in tokens")
|
|
parser.add_argument("--chunk-overlap", type=int, default=200, help="Chunk overlap in tokens")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Get OpenAI API key
|
|
api_key = os.getenv("OPENAI_API_KEY")
|
|
if not api_key:
|
|
print("Error: OPENAI_API_KEY environment variable not set")
|
|
print("Please set your OpenAI API key in the .env file or as an environment variable")
|
|
sys.exit(1)
|
|
|
|
# Validate models
|
|
if not ModelConfig.validate_models(args.embedding_model, args.chat_model):
|
|
print("Error: Invalid model selection")
|
|
sys.exit(1)
|
|
|
|
# Create configuration
|
|
config = RAGConfig(
|
|
openai_api_key=api_key,
|
|
embedding_model=args.embedding_model,
|
|
chat_model=args.chat_model,
|
|
documents_dir=args.docs_dir,
|
|
index_dir=args.index_dir,
|
|
top_k=args.top_k,
|
|
chunk_size=args.chunk_size,
|
|
chunk_overlap=args.chunk_overlap
|
|
)
|
|
|
|
# Initialize RAG system
|
|
rag_system = RAGSystem(config)
|
|
|
|
# Build or load index
|
|
rag_system.build_index(force_rebuild=args.rebuild_index)
|
|
|
|
# Process query or run interactive mode
|
|
if args.question:
|
|
answer = rag_system.query(args.question)
|
|
print(f"Question: {args.question}")
|
|
print(f"Answer: {answer}")
|
|
elif args.interactive:
|
|
rag_system.interactive_mode()
|
|
else:
|
|
print("Please provide either --question or --interactive flag")
|
|
print("Use --help for more information")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |