Spaces:
Running
Running
File size: 8,766 Bytes
c05de71 |
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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 |
"""
Core chatbot implementation for document question answering.
"""
import logging
import os
import time
from typing import Optional, Dict, Any, List
from tqdm import tqdm
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings, StorageContext, load_index_from_storage
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
import config
# Configure logging
logger = logging.getLogger(__name__)
class Chatbot:
"""Chatbot for document question answering using LlamaIndex."""
def __init__(self, config_dict: Optional[Dict[str, Any]] = None):
"""Initialize the chatbot with configuration.
Args:
config_dict: Optional configuration dictionary. If not provided,
configuration is loaded from environment variables.
"""
# Set up basic variables and load configuration
self.config = config_dict or config.get_chatbot_config()
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.
Returns:
API key as string
Raises:
ValueError: If API key is not found
"""
api_key = config.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.
Sets up the LLM, embedding model, and other settings.
Raises:
Exception: If component initialization fails
"""
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", config.LLM_MODEL),
temperature=self.config.get("temperature", config.LLM_TEMPERATURE),
max_tokens=self.config.get("max_tokens", config.LLM_MAX_TOKENS)
)
# 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", config.EMBEDDING_MODEL),
device=self.config.get("device", config.EMBEDDING_DEVICE),
embed_batch_size=self.config.get("embed_batch_size", config.EMBEDDING_BATCH_SIZE)
)
# 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", config.CHUNK_SIZE),
chunk_overlap=self.config.get("chunk_overlap", config.CHUNK_OVERLAP),
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 = None) -> List:
"""Load documents with retry logic.
Args:
data_dir: Directory containing documents to load. If None, uses default.
Returns:
List of loaded documents
Raises:
Exception: If document loading fails after retries
"""
# Try to load documents up to 3 times if there's an error
max_retries = 3
retry_delay = 1
data_dir = data_dir or config.DATA_DIR
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, index_dir: str = None):
"""Create index with error handling.
Args:
documents: Documents to index
index_dir: Directory to store the index. If None, uses default.
Raises:
Exception: If index creation fails
"""
index_dir = index_dir or config.INDEX_DIR
try:
# Check if index already exists
if os.path.exists(os.path.join(index_dir, "index_store.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.
Sets up the query engine from the index.
Raises:
Exception: If query engine initialization fails
"""
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.
Args:
query_text: The question to answer
Returns:
Response as string
Raises:
Exception: If query fails after 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}")
response = self.query_engine.query(query_text)
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.
Performs any necessary cleanup operations.
"""
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}") |