main.py

main.py

@@ -0,0 +1,765 @@
+import os
+import re
+import requests
+import sqlite3
+import threading
+import time
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from sentence_transformers import SentenceTransformer
+from tenacity import retry, stop_after_attempt, wait_exponential
+import mwparserfromhell
+import logging
+import chromadb
+from collections import deque
+from huggingface_hub import login
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+from fastapi.concurrency import run_in_threadpool # Import for running sync code in async app
+
+# --- Configuration ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+CACHE_DIR = "one_piece_cache"
+DB_PATH = os.path.join(CACHE_DIR, "one_piece_data.db")
+CHROMA_DB_PATH = os.path.join(CACHE_DIR, "chroma_db")
+
+# Model Selection
+LLM_MODEL = "google/gemma-2-2b-it"
+EMBED_MODEL = "intfloat/e5-small-v2"
+
+# Hugging Face token
+# It's recommended to pass this via environment variables in production
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Key One Piece categories to crawl
+WIKI_CATEGORIES = {
+    "Characters": ["Straw_Hat_Pirates", "Marines", "Yonko", "Seven_Warlords", "Worst_Generation"],
+    "Devil_Fruits": ["Paramecia", "Zoan", "Logia"],
+    "Locations": ["Islands", "Seas", "Grand_Line", "New_World"],
+    "Story": ["Story_Arcs", "Sagas", "Events"],
+    "Organizations": ["Pirates", "Crews", "Marines", "World_Government"],
+    "Concepts": ["Haki", "Void_Century", "Ancient_Weapons", "Will_of_D"]
+}
+
+# Most important pages (prioritized)
+CRUCIAL_PAGES = [
+    "Monkey_D._Luffy", "Straw_Hat_Pirates", "One_Piece_(Manga)", "Eiichiro_Oda",
+    "Devil_Fruit", "Haki", "Void_Century", "Gol_D._Roger", "Marines", "Yonko",
+    "World_Government", "Grand_Line", "New_World", "One_Piece", "Will_of_D",
+    "Poneglyphs", "Ancient_Weapons", "Roger_Pirates", "God_Valley_Incident",
+    "Joy_Boy", "Sun_God_Nika", "Laugh_Tale", "Rocks_Pirates", "Revolutionary_Army",
+    "Hito_Hito_no_Mi,_Model:_Nika", "Gomu_Gomu_no_Mi", "Five_Elders", "Im",
+    "Marshall_D._Teach", "Blackbeard_Pirates", "Gura_Gura_no_Mi", "Yami_Yami_no_Mi"
+]
+
+# Processing Parameters
+CHUNK_SIZE_TOKENS = 300
+CHUNK_OVERLAP = 2
+MAX_CONTEXT_CHUNKS = 10 # Increased from 8 to include more context
+SIMILARITY_THRESHOLD = 0.35
+REFRESH_INTERVAL = 7 * 24 * 3600 # Weekly refresh
+CONVERSATION_HISTORY_LENGTH = 6 # Enhanced for better context
+
+class OnePieceChatbot:
+    def __init__(self):
+        os.makedirs(CACHE_DIR, exist_ok=True)
+
+        # Login to Hugging Face
+        if HF_TOKEN:
+            try:
+                login(token=HF_TOKEN)
+                logging.info("Successfully logged into Hugging Face.")
+            except Exception as e:
+                 logging.warning(f"Hugging Face login failed: {e}. Proceeding without explicit login.")
+
+
+        # Initialize database and vector store
+        self.db_conn = self._init_db()
+        self.chroma_client, self.chroma_collection = self._init_chroma()
+
+        # Initialize threading control
+        self.data_lock = threading.Lock()
+        self.processing_pages = set()
+        self.initial_processing_done = threading.Event() # Event to signal initial data load
+
+        # Initialize models - These are memory intensive, loaded once
+        try:
+            self.embedder = SentenceTransformer(EMBED_MODEL)
+            logging.info(f"Loaded SentenceTransformer model: {EMBED_MODEL}")
+            self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
+            logging.info(f"Loaded Tokenizer: {LLM_MODEL}")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                LLM_MODEL,
+                device_map="auto", # Automatically select device (GPU if available)
+                torch_dtype=torch.bfloat16 # Use bfloat16 for efficiency
+            )
+            logging.info(f"Loaded LLM Model: {LLM_MODEL}")
+
+            # Initialize text generation pipeline
+            self.generator = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                max_new_tokens=500, # Increased from 300 to avoid cut-off answers
+                temperature=0.2,
+                do_sample=True, # Set to True to match the temperature setting
+                repetition_penalty=1.2 # Added to reduce repetition
+            )
+            logging.info("Initialized text generation pipeline.")
+
+        except Exception as e:
+            logging.error(f"Failed to load models: {e}")
+            # Depending on requirements, you might raise the exception or handle it gracefully
+            # For deployment, you might want to ensure models load or the app fails early
+            raise SystemExit("Failed to load models, exiting.") from e
+
+        # Conversation memory (per instance, this will be global for the API)
+        # A more complex app might manage history per user session
+        self.conversation_history = deque(maxlen=CONVERSATION_HISTORY_LENGTH)
+
+        # Start background data processing thread
+        logging.info("Starting background data processing thread...")
+        thread = threading.Thread(target=self._process_wiki_data, daemon=True)
+        thread.start()
+        logging.info("Background data processing thread started.")
+
+
+    def _init_db(self):
+        """Initialize SQLite database with optimized schema."""
+        conn = sqlite3.connect(DB_PATH, check_same_thread=False) # check_same_thread=False is needed for FastAPI's async nature
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS wiki_data (
+                title TEXT PRIMARY KEY,
+                content TEXT,
+                category TEXT,
+                last_fetched REAL,
+                page_links TEXT
+            )
+        """)
+        conn.execute("CREATE INDEX IF NOT EXISTS idx_category ON wiki_data (category)")
+        conn.commit()
+        logging.info(f"SQLite database initialized at {DB_PATH}")
+        return conn
+
+    def _init_chroma(self):
+        """Initialize ChromaDB for vector storage."""
+        client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
+        collection_name = "one_piece_knowledge"
+        try:
+             collection = client.get_collection(name=collection_name)
+             logging.info(f"Connected to existing ChromaDB collection: {collection_name}")
+        except: # Catch exception if collection doesn't exist
+             collection = client.create_collection(
+                name=collection_name,
+                metadata={"hnsw:space": "cosine"} # Optimize for semantic similarity
+            )
+             logging.info(f"Created new ChromaDB collection: {collection_name}")
+
+        logging.info(f"ChromaDB initialized at {CHROMA_DB_PATH}")
+        return client, collection
+
+    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
+    def _fetch_wiki_page(self, title):
+        """Fetch a page from the One Piece wiki with improved parsing."""
+        logging.debug(f"Attempting to fetch wiki page: {title}")
+        url = f"https://onepiece.fandom.com/api.php?action=parse&page={title}&format=json&prop=wikitext|categories"
+        try:
+            response = requests.get(url, timeout=15) # Increased timeout slightly
+            response.raise_for_status() # Raise an HTTPError for bad responses (4xx or 5xx)
+            data = response.json()
+
+            if "parse" not in data:
+                 logging.warning(f"Could not parse wiki data for {title}")
+                 return None, [], None
+
+            wikitext = data["parse"]["wikitext"]["*"]
+            parsed = mwparserfromhell.parse(wikitext)
+
+            # Clean and extract text - Keep templates for now, focus on removing infoboxes/sidebars potentially
+            # A more robust parser might be needed for complex wiki pages
+            # Simple approach: remove large templates that aren't core content
+            for node in parsed.ifilter_templates():
+                 # Example: remove templates starting with 'Infobox' or 'Character Infobox'
+                 template_name = str(node.name).strip().lower()
+                 if template_name.startswith('infobox') or 'sidebar' in template_name:
+                     try:
+                         parsed.remove(node)
+                     except ValueError:
+                         pass # Node might already be removed
+
+            # Extract internal links
+            links = []
+            for link in parsed.ifilter_wikilinks():
+                link_title = str(link.title).split("#")[0].strip() # Remove section links
+                if ":" not in link_title and len(link_title) > 1 and not link_title.startswith(('File:', 'Category:', 'Template:')):
+                    links.append(link_title)
+
+            # Determine category
+            category = "Other"
+            if "categories" in data["parse"]:
+                categories = [cat["*"] for cat in data["parse"]["categories"]]
+                for cat_type, cat_list in WIKI_CATEGORIES.items():
+                    if any(cat.replace(' ', '_') in [c.replace(' ', '_') for c in categories] for cat in cat_list):
+                        category = cat_type
+                        break
+
+            text = parsed.strip_code().strip() # Strip all wiki code
+            text = re.sub(r'https?://\S+', '', text) # Remove URLs
+            text = re.sub(r'\[\[[^\]]+\]\]', '', text) # Remove any remaining links/categories in text
+            text = re.sub(r'\s+', ' ', text).strip() # Normalize all whitespace to single spaces
+            text = re.sub(r'\n{2,}', '\n\n', text) # Normalize multiple newlines to double newline
+
+            logging.debug(f"Successfully fetched and parsed {title}, category: {category}, links found: {len(links)}")
+            return text, links, category
+
+        except requests.exceptions.RequestException as e:
+            logging.error(f"Request error fetching {title}: {e}")
+            raise # Re-raise to trigger retry
+        except Exception as e:
+            logging.error(f"Error processing wiki data for {title}: {e}")
+            return None, [], None # Return None on other errors
+
+    def _fetch_category_pages(self, category):
+        """Fetch all pages in a specific category."""
+        logging.info(f"Fetching pages for category: {category}")
+        url = f"https://onepiece.fandom.com/api.php?action=query&list=categorymembers&cmtitle=Category:{category}&cmlimit=500&format=json"
+        try:
+            response = requests.get(url, timeout=20) # Increased timeout
+            response.raise_for_status()
+            data = response.json()
+
+            pages = []
+            if "query" in data and "categorymembers" in data["query"]:
+                for member in data["query"]["categorymembers"]:
+                     # Filter out specific namespaces like "File:", "Category:", etc.
+                    if member["ns"] == 0 and "title" in member:
+                         pages.append(member["title"])
+
+            logging.info(f"Found {len(pages)} pages in category {category}")
+            return pages
+        except requests.exceptions.RequestException as e:
+            logging.error(f"Request error fetching category {category}: {e}")
+            return []
+        except Exception as e:
+            logging.error(f"Error processing category {category} members: {e}")
+            return []
+
+
+    def _process_wiki_data(self):
+        """Process wiki pages and store in database and vector store."""
+        logging.info("Background processing: Starting data collection...")
+
+        processed_count = 0
+
+        # First process crucial pages
+        logging.info("Processing crucial pages...")
+        for page in CRUCIAL_PAGES:
+             # Check if already processed recently
+             cur = self.db_conn.execute("SELECT last_fetched FROM wiki_data WHERE title = ?", (page,))
+             result = cur.fetchone()
+             if result and time.time() - result[0] < REFRESH_INTERVAL:
+                 logging.debug(f"Skipping recent crucial page: {page}")
+                 processed_count += 1
+                 continue
+
+             if self._process_page(page):
+                 processed_count += 1
+
+
+        # Then fetch and process category pages
+        logging.info("Processing category pages...")
+        crawled_pages_from_categories = set()
+        for category_type, categories in WIKI_CATEGORIES.items():
+            for category in categories:
+                try:
+                    pages = self._fetch_category_pages(category)
+                    for page in pages:
+                         # Avoid processing pages already processed or in queue from crucial list
+                         if page in CRUCIAL_PAGES: continue
+                         if page in crawled_pages_from_categories: continue
+
+                         crawled_pages_from_categories.add(page) # Track pages found in categories
+
+                         # Check if already processed recently
+                         cur = self.db_conn.execute("SELECT last_fetched FROM wiki_data WHERE title = ?", (page,))
+                         result = cur.fetchone()
+                         if result and time.time() - result[0] < REFRESH_INTERVAL:
+                             logging.debug(f"Skipping recent category page: {page}")
+                             processed_count += 1
+                             continue
+
+                         if self._process_page(page):
+                             processed_count += 1
+
+
+                except Exception as e:
+                    logging.error(f"Error processing category {category}: {e}")
+
+        self.initial_processing_done.set() # Signal that initial processing is complete
+        logging.info(f"Initial data processing finished. Processed {processed_count} pages.")
+        logging.info(f"Vector collection count after initial processing: {self.chroma_collection.count()}")
+
+
+        # Periodically refresh data
+        while True:
+            time.sleep(REFRESH_INTERVAL)
+            logging.info("Starting periodic refresh cycle...")
+
+            # Refresh data for a batch of existing pages ordered by oldest fetch time
+            cur = self.db_conn.execute("SELECT title FROM wiki_data ORDER BY last_fetched ASC LIMIT 200") # Increased batch size
+            pages_to_refresh = [row[0] for row in cur.fetchall()]
+
+            logging.info(f"Refreshing {len(pages_to_refresh)} pages.")
+            for page in pages_to_refresh:
+                 self._process_page(page) # This checks for recent fetch time internally
+
+            logging.info("Periodic refresh cycle finished.")
+
+
+    def _process_page(self, title):
+        """Process a single wiki page."""
+        with self.data_lock: # Use lock to prevent concurrent processing of the same page
+            if title in self.processing_pages:
+                logging.debug(f"Page {title} already in processing queue.")
+                return False # Indicate that processing was skipped
+            # Check if we *really* need to fetch this page again (double check under lock)
+            cur = self.db_conn.execute("SELECT last_fetched FROM wiki_data WHERE title = ?", (title,))
+            result = cur.fetchone()
+            if result and time.time() - result[0] < REFRESH_INTERVAL:
+                 logging.debug(f"Page {title} already processed recently under lock.")
+                 return False # Indicate that processing was skipped
+
+            self.processing_pages.add(title)
+            logging.info(f"Processing page: {title}")
+
+        try:
+            # Fetch and process the page - happens outside the initial lock
+            content, links, category = self._fetch_wiki_page(title)
+            if not content:
+                logging.warning(f"No content fetched for {title}")
+                return False # Indicate failure
+
+            # Store in SQLite (requires lock)
+            with self.data_lock:
+                 self.db_conn.execute(
+                    "INSERT OR REPLACE INTO wiki_data VALUES (?, ?, ?, ?, ?)",
+                    (title, content, category, time.time(), ','.join(links))
+                 )
+                 self.db_conn.commit()
+                 logging.debug(f"Stored {title} in SQLite.")
+
+            # Chunk and embed
+            chunks = self._chunk_text(content, title)
+            if chunks:
+                try:
+                    # ChromaDB operations can be done outside the main data_lock
+                    # but ensure ChromaDB client/collection access is thread-safe if needed,
+                    # although PersistentClient is generally safe.
+                    embeddings = self.embedder.encode(chunks, convert_to_tensor=False).tolist()
+                    ids = [f"{title}::{i}" for i in range(len(chunks))]
+                    metadatas = [{"source": title, "category": category} for _ in chunks]
+
+                    # Remove old chunks for this page before upserting new ones
+                    try:
+                        old_ids = self.chroma_collection.get(where={"source": title}, include=[])["ids"]
+                        if old_ids:
+                            self.chroma_collection.delete(ids=old_ids)
+                            logging.debug(f"Removed {len(old_ids)} old chunks for {title} from ChromaDB.")
+                    except Exception as delete_e:
+                         logging.warning(f"Could not delete old chunks for {title} from ChromaDB: {delete_e}")
+
+
+                    self.chroma_collection.upsert(
+                        ids=ids,
+                        embeddings=embeddings,
+                        documents=chunks,
+                        metadatas=metadatas
+                    )
+                    logging.info(f"Processed {title}: {len(chunks)} chunks added to ChromaDB.")
+                except Exception as chroma_e:
+                     logging.error(f"Error adding/updating chunks for {title} in ChromaDB: {chroma_e}")
+                     # Decide if you want to re-raise or just log the error
+                     return False # Indicate failure if vector store update fails
+
+            # Process linked pages if needed (can be done outside the main lock)
+            # Add linked pages to processing queue
+            if links:
+                logging.debug(f"Adding linked pages from {title} to processing queue.")
+                for link in links[:10]: # Limit to first 10 links to avoid too much sprawl
+                     # Queue the page processing, don't block here
+                     threading.Thread(target=self._process_page, args=(link,), daemon=True).start()
+
+
+            return True # Indicate success
+
+        except Exception as e:
+            logging.error(f"Caught unexpected error during processing of {title}: {e}")
+            return False # Indicate failure
+        finally:
+            # Ensure the page is removed from processing set even on error
+            with self.data_lock:
+                 if title in self.processing_pages:
+                    self.processing_pages.remove(title)
+                    logging.debug(f"Removed {title} from processing queue.")
+
+
+    def _chunk_text(self, text, title):
+        """Split text into chunks for embedding with sentence awareness."""
+        # Use nltk for better sentence tokenization if possible, but regex is ok for a start
+        sentences = re.split(r'(?<=[.!?])\s+', text)
+        chunks, current_sentences = [], []
+        current_chunk_tokens = 0
+
+        # Simple check for minimal content
+        if len(sentences) < 2 and len(text.split()) < 50:
+             logging.debug(f"Page {title} is too short for chunking ({len(text.split())} words). Skipping.")
+             return []
+
+
+        for i, sentence in enumerate(sentences):
+            # Use tokenizer for token counting
+            sentence_tokens = len(self.tokenizer.encode(sentence, add_special_tokens=False))
+
+            # Check if adding the sentence would exceed chunk size OR if it's the last sentence
+            # If we add the sentence, what would the new token count be including a space?
+            # (approximate)
+            new_token_count = current_chunk_tokens + sentence_tokens + (1 if current_chunk_tokens > 0 else 0)
+
+            if new_token_count > CHUNK_SIZE_TOKENS and current_sentences:
+                # If adding sentence exceeds chunk size, save current sentences as a chunk
+                chunk_text = " ".join(current_sentences).strip()
+                if chunk_text: # Ensure chunk is not empty
+                    chunks.append(chunk_text)
+                logging.debug(f"Chunked: {len(chunk_text.split())} words, {len(self.tokenizer.encode(chunk_text, add_special_tokens=False))} tokens")
+
+                # Start new chunk with overlap
+                overlap_sentences = current_sentences[-CHUNK_OVERLAP:] if len(current_sentences) > CHUNK_OVERLAP else current_sentences
+                current_sentences = overlap_sentences
+                current_chunk_tokens = sum(len(self.tokenizer.encode(s, add_special_tokens=False)) for s in current_sentences)
+
+            current_sentences.append(sentence)
+            current_chunk_tokens += sentence_tokens + (1 if current_chunk_tokens > 0 else 0) # Add space token count
+
+        # Add the last chunk if any sentences remain
+        if current_sentences:
+            chunk_text = " ".join(current_sentences).strip()
+            if chunk_text: # Ensure chunk is not empty
+                 chunks.append(chunk_text)
+            logging.debug(f"Final Chunked: {len(chunk_text.split())} words, {len(self.tokenizer.encode(chunk_text, add_special_tokens=False))} tokens")
+
+        # Filter out chunks that are too short
+        return [chunk for chunk in chunks if len(chunk.split()) > 20]
+
+
+    def _interpret_query(self, query):
+        """Interpret and expand the query using conversation history."""
+        # Check if initial data processing is done before using history for complex interpretation
+        if not self.initial_processing_done.is_set() or len(query.split()) > 3 and not (query.lower().startswith("and ") or query.lower().startswith("what about ")):
+             # If initial data not ready, or query is already substantial, return original
+             return query
+
+        # Handle follow-up questions and vague queries using the LLM
+        logging.debug(f"Interpreting query: '{query}'")
+        try:
+            prompt = f"""
+Based on this conversation history:
+{self._format_history()}
+
+The user asked a question: "{query}"
+Please interpret this as a complete, standalone question about One Piece, incorporating context from the history if necessary. Ensure the reformulated question is clear and specific, even if the original query was vague or a follow-up.
+Only provide the complete reformulated question and nothing else.
+"""
+            # Use the generator for interpretation
+            # Set max_new_tokens low as we only expect a short question
+            interpretation_response = self.generator(
+                 prompt,
+                 max_new_tokens=50,
+                 temperature=0.5,
+                 do_sample=True,
+                 repetition_penalty=1.1
+            )[0]["generated_text"]
+
+            # Extract the reformulated question (look for the part after the instruction)
+            # This is a bit fragile, depends on LLM output format
+            if "Only provide the complete reformulated question and nothing else:" in interpretation_response:
+                 interpreted_query = interpretation_response.split("Only provide the complete reformulated question and nothing else:")[-1].strip()
+            else:
+                 # Fallback if the LLM doesn't follow the format strictly
+                 # Try to clean up surrounding text
+                 lines = interpreted_query.split('\n')
+                 interpreted_query = lines[-1].strip() if lines else interpreted_query.strip()
+
+
+            # Basic cleaning
+            interpreted_query = interpreted_query.replace('"', '').strip()
+            logging.info(f"Interpreted '{query}' as '{interpreted_query}'")
+            return interpreted_query
+
+        except Exception as e:
+            logging.error(f"Error interpreting query '{query}': {e}. Using original query.")
+            return query # Fallback to original query on error
+
+
+    def _find_relevant_chunks(self, query):
+        """Find relevant chunks using vector similarity."""
+        # Wait briefly for initial data if not ready. Avoid long blocking.
+        if not self.initial_processing_done.is_set():
+             logging.warning("Initial data processing not complete. Search results may be limited.")
+             # Give it a few seconds just in case it's almost done, but don't block indefinitely
+             self.initial_processing_done.wait(timeout=5) # Wait up to 5 seconds
+
+        interpreted_query = self._interpret_query(query)
+
+        # Enhanced query expansion for key topics - performed *after* interpretation
+        # This adds specific keywords to the query that might help retrieve relevant documents
+        keywords_to_add = []
+        lower_query = interpreted_query.lower()
+
+        if "joy boy" in lower_query or "nika" in lower_query:
+             keywords_to_add.extend(["Hito Hito no Mi Model Nika", "Sun God Nika"])
+        if "blackbeard" in lower_query and "devil fruit" in lower_query:
+             keywords_to_add.extend(["multiple devil fruits", "Yami Yami no Mi", "Gura Gura no Mi"])
+        if "gorosei" in lower_query or "im" in lower_query:
+             keywords_to_add.extend(["Five Elders", "Empty Throne", "World Government"])
+        if "void century" in lower_query:
+             keywords_to_add.extend(["Poneglyphs", "Ancient Kingdom", "Ohara"])
+
+
+        if keywords_to_add:
+            interpreted_query_with_keywords = interpreted_query + " " + " ".join(keywords_to_add)
+            logging.debug(f"Expanded query: {interpreted_query_with_keywords}")
+        else:
+            interpreted_query_with_keywords = interpreted_query
+
+
+        # Search vector database
+        try:
+            query_embedding = self.embedder.encode(interpreted_query_with_keywords).tolist()
+            results = self.chroma_collection.query(
+                query_embeddings=[query_embedding],
+                n_results=MAX_CONTEXT_CHUNKS,
+                include=["documents", "metadatas", "distances"]
+            )
+        except Exception as e:
+            logging.error(f"Error querying ChromaDB: {e}")
+            return [], [] # Return empty if search fails
+
+
+        chunks = []
+        sources = set() # Use a set to keep sources unique
+
+        if results and results["documents"]:
+             # Only include chunks above similarity threshold
+             for i, doc in enumerate(results["documents"][0]):
+                 distance = results["distances"][0][i]
+                 similarity = 1 - distance # Cosine distance is 1 - cosine similarity
+
+                 logging.debug(f"Chunk {i+1}: Source={results['metadatas'][0][i]['source']}, Distance={distance:.4f}, Similarity={similarity:.4f}")
+
+                 if similarity >= SIMILARITY_THRESHOLD:
+                     chunks.append(doc)
+                     sources.add(results["metadatas"][0][i]["source"])
+
+        logging.info(f"Found {len(chunks)} relevant chunks for query '{query}'. Sources: {list(sources)}")
+        return chunks, list(sources)
+
+
+    def _extract_entities(self, text):
+        """Extract potential One Piece entities from text."""
+        # This method isn't currently used in answer_question, but kept from original.
+        # Simple heuristic: look for capitalized words that might be names/places
+        entities = re.findall(r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b', text)
+        # Filter out common English words and short words that aren't likely entities
+        common_words = {"The", "A", "An", "In", "On", "At", "For", "With", "And", "But", "Or", "Not", "Is", "Are", "Was", "Were"}
+        return [e for e in entities if len(e) > 3 and e not in common_words]
+
+
+    def _format_history(self):
+        """Format conversation history for the LLM prompt."""
+        if not self.conversation_history:
+            return "No recent conversation history."
+
+        history = "Recent conversation history:\n"
+        for i, (q, a) in enumerate(self.conversation_history):
+            history += f"Turn {i+1}:\nUser: {q}\nAssistant: {a}\n"
+        return history
+
+    def answer_question(self, question: str):
+        """Answer a question using retrieved context and conversation history."""
+        logging.info(f"Received question: '{question}'")
+
+        # Wait for initial data processing if not finished.
+        # This makes the first few requests potentially slower but ensures some data is loaded.
+        # Consider a proper "loading" status endpoint if this is too slow.
+        if not self.initial_processing_done.is_set():
+            logging.warning("Initial data processing not yet complete. Waiting up to 60s...")
+            if not self.initial_processing_done.wait(timeout=60):
+                logging.error("Initial data processing timed out. Cannot answer reliably.")
+                return "The knowledge base is still loading. Please try again in a few minutes."
+            else:
+                logging.info("Initial data processing finished while waiting.")
+
+
+        chunks, sources = self._find_relevant_chunks(question)
+
+        if not chunks:
+             logging.warning(f"No relevant chunks found for question: '{question}'")
+             # Add a fallback or a "I don't know" response if no context is found
+             fallback_prompt = f"""You are an expert on the One Piece manga and anime. The user asked: "{question}". However, no relevant specific information was found in your knowledge base. Provide a general, helpful answer based on your broad understanding of One Piece, or state that you don't have specific information on this topic. Do not invent facts.
+
+IMPORTANT: Start immediately with your answer."""
+             try:
+                 response = self.generator(fallback_prompt, max_new_tokens=200, temperature=0.7, do_sample=True)[0]["generated_text"].strip()
+                 # Clean up prompt instructions
+                 response = re.sub(r'^.*?IMPORTANT: Start immediately with your answer\.', '', response, flags=re.DOTALL).strip()
+                 answer = response
+                 # Add to history (can decide if you want "I don't know" in history)
+                 self.conversation_history.append((question, answer))
+                 return answer
+             except Exception as e:
+                  logging.error(f"Error generating fallback response: {e}")
+                  answer = "I couldn't find specific information about that in my knowledge base."
+                  self.conversation_history.append((question, answer))
+                  return answer
+
+
+        # Construct the prompt for the LLM
+        # Enhanced prompt with clear output instructions
+        prompt = f"""You are an expert on the One Piece manga and anime. Answer the following question based *only* on the provided context and your knowledge of One Piece lore.
+
+{self._format_history()}
+
+Context information:
+{chr(10).join(chunks)}
+
+Question: {question}
+
+Provide a detailed, accurate answer based on the context above. If the context doesn't contain enough information to fully answer, use your general One Piece knowledge but prioritize information from the context. Explain connections between characters and events clearly. Structure your answer logically.
+
+IMPORTANT: Your answer must be directly useful and not include phrases like "based on the context" or "answer the question". Start immediately with your answer. Ensure your answer is cohesive and well-formatted.
+"""
+
+        logging.debug(f"Sending prompt to LLM:\n{prompt}")
+
+        try:
+            # Get the response from the generator pipeline
+            response = self.generator(prompt)[0]["generated_text"]
+            logging.debug(f"Raw LLM response:\n{response}")
+
+            # Extract just the answer portion and clean it
+            # Use a pattern that matches the instruction part robustly
+            answer_match = re.search(r'IMPORTANT:.*?Start immediately with your answer\.(.*)', response, re.DOTALL)
+            if answer_match:
+                answer = answer_match.group(1).strip()
+            else:
+                # Fallback extraction if the instruction pattern is not found
+                # Attempt to find the question again and take everything after it
+                answer_parts = response.split("Question: " + question)
+                if len(answer_parts) > 1:
+                     answer = answer_parts[-1].strip()
+                else:
+                     # If still not found, take the last significant block or the whole thing
+                     logging.warning("Could not find extraction pattern, falling back to end of response.")
+                     answer = response.strip() # Or implement more sophisticated fallback logic
+
+
+            # Clean up any lingering prompt instructions or metadata accidentally generated
+            answer = re.sub(r'^(.*?)(?:IMPORTANT:|Based on this conversation history:|Context information:|Question:)', '', answer, flags=re.DOTALL | re.IGNORECASE).strip()
+            answer = re.sub(r'\s*Sources:\s*.*$', '', answer, flags=re.DOTALL) # Remove any "Sources:" line the LLM might invent
+
+
+        except Exception as e:
+            logging.error(f"Error generating response from LLM: {e}")
+            answer = "Sorry, I encountered an error while generating the response."
+
+
+        # Add to conversation history (user question and generated answer)
+        self.conversation_history.append((question, answer))
+        logging.debug(f"Added to history: Q='{question}', A='{answer[:50]}...'")
+
+        # Format response with better source attribution
+        if sources:
+            # Ensure sources are clean titles, not links or complex text
+            clean_sources = [s.replace('_', ' ') for s in sources] # Simple cleaning
+            sources_list = list(clean_sources)[:5] # Limit to top 5 sources for brevity in output
+            sources_str = ", ".join(sources_list)
+            # Append sources clearly
+            return f"{answer}\n\nSources: {sources_str}"
+
+        return answer
+
+# --- FastAPI Application Setup ---
+
+# Initialize the chatbot instance globally
+# This happens when the script is first imported/run, loading models and starting threads
+try:
+    chatbot = OnePieceChatbot()
+    logging.info("OnePieceChatbot instance created.")
+except Exception as e:
+    logging.critical(f"Failed to initialize OnePieceChatbot: {e}")
+    # In a real deployment, this might stop the container or signal an error
+    # For now, we'll let the app start but requests will likely fail
+
+app = FastAPI()
+
+# Pydantic model for the request body
+class QuestionRequest(BaseModel):
+    question: str
+
+# Root endpoint
+@app.get("/")
+async def read_root():
+    return {"message": "One Piece Chatbot API is running. Send a POST request to /ask with your question."}
+
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    # Add more sophisticated checks here if needed (e.g., db connection, model loaded)
+    if not hasattr(chatbot, 'generator') or chatbot.generator is None:
+         return {"status": "error", "message": "LLM model not loaded"}, 500
+
+    # Check if the background thread has completed initial processing
+    if not chatbot.initial_processing_done.is_set():
+         return {"status": "warning", "message": "Initial data processing still in progress. Some answers may be limited."}, 200 # Or 503 Service Unavailable
+
+    # Optional: Check ChromaDB count to ensure it's not empty after initialization
+    try:
+         count = chatbot.chroma_collection.count()
+         if count == 0:
+             return {"status": "warning", "message": "Knowledge base is empty after initialization. Data fetching might have failed."}, 200
+         return {"status": "ok", "message": "Chatbot is ready.", "knowledge_base_size": count}, 200
+    except Exception as e:
+         logging.error(f"Health check failed during ChromaDB count: {e}")
+         return {"status": "warning", "message": f"Health check encountered an issue: {e}"}, 200
+
+
+# Endpoint to answer questions
+@app.post("/ask")
+async def ask_question_endpoint(request: QuestionRequest):
+    """
+    Submit a question about One Piece and get an answer.
+    """
+    question = request.question
+    if not question or not question.strip():
+        return {"answer": "Please provide a question."}
+
+    # Run the synchronous answer_question method in a thread pool
+    # This prevents blocking the FastAPI event loop
+    try:
+        answer = await run_in_threadpool(chatbot.answer_question, question)
+        return {"answer": answer}
+    except Exception as e:
+        logging.error(f"Error processing question '{question}': {e}")
+        return {"answer": "Sorry, an internal error occurred while processing your question."}, 500 # Return 500 status code on error
+
+# To run this application:
+# 1. Save the code as a Python file (e.g., main.py).
+# 2. Make sure you have the necessary libraries installed:
+#    pip install fastapi uvicorn transformers sentence-transformers tenacity mwparserfromhell chromadb requests sqlite3 accelerate bitsandbytes torch huggingface_hub
+#    (You might need additional dependencies for torch/accelerate based on your hardware, e.g., CUDA)
+# 3. Run from your terminal: uvicorn main:app --reload
+# For deployment, you'll typically use a production server like Gunicorn with Uvicorn workers:
+#    gunicorn -w 4 -k uvicorn.workers.UvicornWorker main:app
+# Ensure the HF_TOKEN environment variable is set in your deployment environment.