from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings from llama_index.llms.anthropic import Anthropic from llama_index.embeddings.huggingface import HuggingFaceEmbedding from llama_index.core.node_parser import SentenceSplitter from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler from llama_index.core import StorageContext, load_index_from_storage import logging import os from dotenv import load_dotenv import time from typing import Optional, Dict, Any from tqdm import tqdm # Set up logging to track what the chatbot is doing logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) # Disable tokenizer parallelism warnings os.environ["TOKENIZERS_PARALLELISM"] = "false" # Create a directory for storing the index INDEX_DIR = "index" if not os.path.exists(INDEX_DIR): os.makedirs(INDEX_DIR) class Chatbot: def __init__(self, config: Optional[Dict[str, Any]] = None): """Initialize the chatbot with configuration.""" # Set up basic variables and load configuration self.config = config or {} self.api_key = self._get_api_key() self.index = None self.query_engine = None self.llm = None self.embed_model = None # Set up debugging tools to help track any issues self.debug_handler = LlamaDebugHandler(print_trace_on_end=True) self.callback_manager = CallbackManager([self.debug_handler]) # Set up all the components needed for the chatbot self._initialize_components() def _get_api_key(self) -> str: """Get API key from environment or config.""" # Load the API key from environment variables or config file load_dotenv() api_key = os.getenv("ANTHROPIC_API_KEY") or self.config.get("api_key") if not api_key: raise ValueError("API key not found in environment or config") return api_key def _initialize_components(self): """Initialize all components with proper error handling.""" try: # Set up the language model (Claude) with our settings logger.info("Setting up Claude language model...") self.llm = Anthropic( api_key=self.api_key, model=self.config.get("model", "claude-3-7-sonnet-20250219"), temperature=self.config.get("temperature", 0.1), max_tokens=self.config.get("max_tokens", 2048) # Allow for longer responses ) # Set up the model that converts text into numbers (embeddings) logger.info("Setting up text embedding model...") self.embed_model = HuggingFaceEmbedding( model_name=self.config.get("embedding_model", "sentence-transformers/all-MiniLM-L6-v2"), device=self.config.get("device", "cpu"), embed_batch_size=self.config.get("embed_batch_size", 8) ) # Configure all the settings for the chatbot logger.info("Configuring chatbot settings...") Settings.embed_model = self.embed_model Settings.text_splitter = SentenceSplitter( chunk_size=self.config.get("chunk_size", 1024), chunk_overlap=self.config.get("chunk_overlap", 100), paragraph_separator="\n\n" ) Settings.llm = self.llm Settings.callback_manager = self.callback_manager logger.info("Components initialized successfully") except Exception as e: logger.error(f"Error initializing components: {e}") raise def load_documents(self, data_dir: str = "data"): """Load documents with retry logic.""" # Try to load documents up to 3 times if there's an error max_retries = 3 retry_delay = 1 for attempt in range(max_retries): try: logger.info(f"Loading documents from {data_dir}...") documents = SimpleDirectoryReader(data_dir).load_data() logger.info(f"Loaded {len(documents)} documents") return documents except Exception as e: if attempt < max_retries - 1: logger.warning(f"Attempt {attempt + 1} failed: {e}. Retrying in {retry_delay} seconds...") time.sleep(retry_delay) else: logger.error(f"Failed to load documents after {max_retries} attempts: {e}") raise def create_index(self, documents): """Create index with error handling.""" try: # Check if index already exists if os.path.exists(os.path.join(INDEX_DIR, "index.json")): logger.info("Loading existing index...") storage_context = StorageContext.from_defaults(persist_dir=INDEX_DIR) self.index = load_index_from_storage(storage_context) logger.info("Index loaded successfully") return # Create a new index if none exists logger.info("Creating new index...") with tqdm(total=1, desc="Creating searchable index") as pbar: self.index = VectorStoreIndex.from_documents(documents) # Save the index self.index.storage_context.persist(persist_dir=INDEX_DIR) pbar.update(1) logger.info("Index created and saved successfully") except Exception as e: logger.error(f"Error creating/loading index: {e}") raise def initialize_query_engine(self): """Initialize query engine with error handling.""" try: # Set up the system that will handle questions logger.info("Initializing query engine...") self.query_engine = self.index.as_query_engine() logger.info("Query engine initialized successfully") except Exception as e: logger.error(f"Error initializing query engine: {e}") raise def query(self, query_text: str) -> str: """Execute a query with error handling and retries.""" # Try to answer questions up to 3 times if there's an error max_retries = 3 retry_delay = 1 for attempt in range(max_retries): try: logger.info(f"Executing query: {query_text}") print("\nThinking...", end="", flush=True) response = self.query_engine.query(query_text) print(" Done!") logger.info("Query executed successfully") return str(response) except Exception as e: if attempt < max_retries - 1: logger.warning(f"Attempt {attempt + 1} failed: {e}. Retrying in {retry_delay} seconds...") time.sleep(retry_delay) else: logger.error(f"Failed to execute query after {max_retries} attempts: {e}") raise def cleanup(self): """Clean up resources.""" try: # Clean up any resources we used logger.info("Cleaning up resources...") logger.info("Cleanup completed successfully") except Exception as e: logger.error(f"Error during cleanup: {e}") def main(): # Set up all the configuration settings for the chatbot config = { "model": "claude-3-7-sonnet-20250219", "temperature": 0.1, "max_tokens": 2048, # Allow for longer responses "embedding_model": "sentence-transformers/all-MiniLM-L6-v2", "device": "cpu", "embed_batch_size": 8, "chunk_size": 1024, "chunk_overlap": 100 } chatbot = None try: # Create and set up the chatbot print("\nInitializing chatbot...") chatbot = Chatbot(config) # Load the documents we want to analyze documents = chatbot.load_documents() # Create a searchable index from the documents chatbot.create_index(documents) # Set up the system that will handle questions chatbot.initialize_query_engine() print("\nChatbot is ready! You can ask questions about your documents.") print("Type 'exit' to quit.") print("-" * 50) while True: # Get user input question = input("\nYour question: ").strip() # Check if user wants to exit if question.lower() in ['exit', 'quit', 'bye']: print("\nGoodbye!") break # Skip empty questions if not question: continue # Get and print the response response = chatbot.query(question) print("\nAnswer:", response) print("-" * 50) except Exception as e: logger.error(f"Error in main: {e}") finally: # Clean up when we're done if chatbot: chatbot.cleanup() if __name__ == "__main__": main()