rag_system.py

import os
import json
import requests
import time
from typing import List, Dict, Any, Optional, Tuple, Iterator
from tqdm import tqdm
import re
from pymilvus import MilvusClient
from sentence_transformers import SentenceTransformer
import hashlib
from pathlib import Path
import glob

from config import *
from temporal_analysis import TemporalAnalyzer
from pattern_discovery import PatternDiscoveryAgent

# Configuration
MILVUS_URI = os.getenv("MILVUS_URI", MILVUS_URI)
COLLECTION_NAME = os.getenv("COLLECTION_NAME", COLLECTION_NAME)
LLM_API_URL = os.getenv("LLM_API_URL", LLM_API_URL)
BRAVE_API_KEY = os.getenv("BRAVE_API_KEY", BRAVE_API_KEY)
FILES_DIR = os.getenv("FILES_DIR", FILES_DIR)
META_DIR = os.getenv("META_DIR", META_DIR)
LLM_MODEL_NAME = os.getenv("LLM_MODEL_NAME", LLM_MODEL_NAME)

WOS_API_URL = os.getenv("WOS_API_URL", WOS_API_URL)
WOS_API_KEY = os.getenv("WOS_API_KEY", WOS_API_KEY)


# Initialize embedding model
embedding_model = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2', device='cpu')


class RAGSystem:
    """
    Enhanced RAG system:

    - Conversation memory
    - Query expansion
    - Chain of thought reasoning
    - Tool use (Brave Search API)
    """

    def __init__(
            self,
            milvus_uri: str = MILVUS_URI,
            collection_name: str = COLLECTION_NAME,
            llm_api_url: str = LLM_API_URL,
            max_tokens: int = 2048,
            temperature: float = 0.7,
            search_limit: int = 5,
            dense_weight: float = 0.6,
            sparse_weight: float = 0.4,
            memory_enabled: bool = True,
            memory_turns: int = 3,
            query_expansion_enabled: bool = True,
            auto_load_metadata: bool = False
    ):
        """Initialize the RAG system"""
        self.milvus_client = MilvusClient(uri=milvus_uri)
        self.collection_name = collection_name

        # LLM configuration
        self.llm_api_url = llm_api_url
        self.max_tokens = max_tokens
        self.temperature = temperature

        # Search parameters
        self.search_limit = search_limit
        self.dense_weight = dense_weight
        self.sparse_weight = sparse_weight

        # Tool use
        self.use_web_search = False
        self.brave_api_key = BRAVE_API_KEY

        # Memory
        self.memory_enabled = memory_enabled
        self.memory_turns = memory_turns

        # Query expansion toggle
        self.query_expansion_enabled = query_expansion_enabled
        self.query_expansion_mode = "conservative"

        self.metadata = {}
        self.metadata_cache = {}

        self.temporal_analyzer = TemporalAnalyzer(self)
        self.pattern_discovery = PatternDiscoveryAgent(self)

        self.files_dir = Path(FILES_DIR) if FILES_DIR else None
        self.meta_dir = Path(META_DIR) if META_DIR else None

        self.use_wos_search = False
        self.wos_api_key = WOS_API_KEY
        self.wos_api_url = WOS_API_URL

        if not self._check_wos_connection():
            print("Warning: Could not connect to Web of Science API. WoS search will not work.")
            self.use_wos_search = False

        # Check if Milvus collection exists
        if not self._check_collection_exists():
            raise ValueError(f"Collection '{collection_name}' does not exist in Milvus")

        if auto_load_metadata and self.files_dir and self.meta_dir:
            try:
                print("Loading document metadata...")
                metadata_start = time.time()
                self.load_all_metadata()
                print(f"Metadata loaded in {time.time() - metadata_start:.2f}s")
            except Exception as e:
                print(f"Error loading metadata: {e}")

        # Check if LLM is accessible
        if not self._check_llm_connection():
            print("Warning: Could not connect to LLM API. Generation will not work.")

    # Check API Endpoints
    def _check_collection_exists(self) -> bool:
        """Check if Milvus collection exists"""
        try:
            collections = self.milvus_client.list_collections()
            return self.collection_name in collections
        except Exception as e:
            print(f"Error checking Milvus collection: {e}")
            return False

    def _check_llm_connection(self) -> bool:
        """Check if LLM API is accessible"""
        try:
            # health check
            response = requests.get(self.llm_api_url.replace("/v1/chat/completions", "/health"))
            return response.status_code == 200
        except Exception as e:
            print(f"Error connecting to LLM API: {e}")
            return False

    def _check_wos_connection(self) -> bool:
        """Check if WOS API accessible"""
        if not self.wos_api_key:
            print("Web of Science API key not provided")
            return False

        try:
            headers = {
                "X-ApiKey": self.wos_api_key,
                "Accept": "application/json"
            }

            # Use documents endpoint for WoS API
            test_url = f"{self.wos_api_url}/documents"

            # proper WoS query syntax with tags
            params = {
                "q": "TS=(science)",
                "limit": 1,
                "page": 1
            }

            print(f"Testing WoS connection with URL: {test_url}")
            print(f"Test query: {params['q']}")

            response = requests.get(
                test_url,
                headers=headers,
                params=params,
                timeout=10
            )

            print(f"WoS API Response Status: {response.status_code}")

            if response.status_code == 200:
                print("Web of Science API connection successful")
                data = response.json()
                total_results = data.get('metadata', {}).get('total', 'unknown')
                print(f"Test query returned {total_results} total results")
                return True
            else:
                print(f"WoS API Response: {response.text}")
                if response.status_code == 401:
                    print("Web of Science API authentication failed - check API key")
                elif response.status_code == 400:
                    print("Web of Science API bad request - query syntax issue")
                elif response.status_code == 403:
                    print("Web of Science API access forbidden - check your subscription")
                return False

        except Exception as e:
            print(f"Error connecting to Web of Science API: {e}")
            return False

    def switch_collection(self, new_collection_name: str, auto_load_metadata: bool = False) -> bool:
        """
        Switch to different Milvus collection
        :new_collection_name: name of collection
        :auto_load_metadata: load metadata immediately or not
        """

        try:
            collections = self.milvus_client.list_collections()
            if new_collection_name not in collections:
                print(f"Collection '{new_collection_name}' does not exist in Milvus")
                return False
        except Exception as e:
            print(f"Error checking collections: {e}")
            return False

        self.collection_name = new_collection_name

        # Load cached metadata
        if new_collection_name in self.metadata_cache:
            self.metadata = self.metadata_cache[new_collection_name]
            print(
                f"Switched to collection '{new_collection_name}' (loaded {len(self.metadata)} cached metadata entries)")
        else:
            self.metadata = {}
            if auto_load_metadata and self.files_dir and self.meta_dir:
                try:
                    print(f"Loading metadata for collection '{new_collection_name}'...")
                    self.load_all_metadata()
                except Exception as e:
                    print(f"Error loading metadata: {e}")
            print(f"Switched to collection '{new_collection_name}'")

        return True

    # Metadata Handling is Below
    def _get_first_value(self, text):
        """Extract first value from semicolon-separated text"""
        if not text:
            return ''
        return text.split(';')[0].strip()

    def _get_first_match(self, metadata, candidates):
        """Get first matching field from list of candidates"""
        for candidate in candidates:
            if candidate in metadata and metadata[candidate]:
                return metadata[candidate]
        return ''

    def _get_publication_date(self, date_text):
        """Enhanced publication date extraction"""
        if not date_text:
            return ''

        # Look for 4-digit years in historical range
        year_match = re.search(r'\b(1[4-8]\d{2})\b', date_text)
        if year_match:
            return year_match.group(1)

        # Handle complex EEBO format
        dates = date_text.split(';')
        for date in dates:
            date = date.strip()
            # Look for dates in brackets like [1586]
            if date.startswith('[') and date.endswith(']'):
                bracket_content = date[1:-1]
                # Extract year from bracket content
                year_match = re.search(r'\b(1[4-8]\d{2})\b', bracket_content)
                if year_match:
                    return year_match.group(1)
                return bracket_content

            # Look for plain 4-digit years in historical range
            year_match = re.search(r'\b(1[4-8]\d{2})\b', date)
            if year_match:
                return year_match.group(1)

        # Fallback to first date if no year pattern found
        return dates[0].strip() if dates else ''

    def _load_metadata(self, file_path):
        """Load metadata for specific file with field mapping"""
        try:
            # Convert text file path to metadata file path
            relative_path = Path(file_path).relative_to(self.files_dir)
            meta_file = self.meta_dir / relative_path

            if meta_file.exists():
                with open(meta_file, 'r', encoding='utf-8') as f:
                    meta_content = f.read().strip()

                # Parse the metadata text file
                metadata = {}
                current_key = None
                current_value = []

                for line in meta_content.split('\n'):
                    if ':' in line:
                        if current_key:
                            metadata[current_key] = '; '.join(current_value)

                        key, value = line.split(':', 1)
                        current_key = key.strip()
                        current_value = [value.strip()]
                    else:
                        if current_key and line.strip():
                            current_value.append(line.strip())

                if current_key:
                    metadata[current_key] = '; '.join(current_value)

                # Enhanced metadata mapping
                clean_metadata = self._map_metadata_fields(metadata)
                return clean_metadata

        except Exception as e:
            print(f"Error loading metadata for {file_path}: {e}")

        return None

    def load_all_metadata(self):
        """Load metadata for all files in directory"""
        if not self.files_dir or not self.meta_dir:
            print("Files directory and metadata directory must be set")
            return

        file_paths = glob.glob(os.path.join(self.files_dir, "**"), recursive=True)
        file_paths = [f for f in file_paths if os.path.isfile(f)]

        print(f"Loading metadata for {len(file_paths)} files...")

        for file_path in tqdm(file_paths):
            filename = os.path.basename(file_path)
            metadata = self._load_metadata(file_path)
            if metadata:
                self.metadata[filename] = metadata

        self.metadata_cache[self.collection_name] = self.metadata.copy()

        print(f"Loaded metadata for {len(self.metadata)} files")

    def _map_metadata_fields(self, metadata):
        clean_metadata = {}

        # Title mapping
        title_candidates = ['TITLE', 'Title', 'title']
        title = self._get_first_match(metadata, title_candidates)
        clean_metadata['title'] = title[:400] + '...' if len(title) > 400 else title

        # Author mapping
        author_candidates = ['AUTHOR', 'Author', 'author', 'Person', 'PERSON', 'person']
        clean_metadata['author'] = self._get_first_match(metadata, author_candidates)

        # Date mapping
        date_candidates = ['DATE', 'Date', 'date']
        date_value = self._get_first_match(metadata, date_candidates)
        clean_metadata['date'] = self._get_publication_date(date_value)

        # Language mapping
        language_candidates = ['LANGUAGE', 'Language', 'language']
        clean_metadata['language'] = self._get_first_match(metadata, language_candidates)

        # Publisher mapping
        publisher_candidates = ['PUBLISHER', 'Publisher', 'publisher']
        clean_metadata['publisher'] = self._get_first_value(self._get_first_match(metadata, publisher_candidates))

        # Publication place mapping
        pubplace_candidates = ['PUBPLACE', 'Pubplace', 'pubplace', 'Publication Place']
        clean_metadata['pubplace'] = self._get_first_value(self._get_first_match(metadata, pubplace_candidates))

        # ID mapping
        id_candidates = ['IDNO', 'ID', 'id', 'Identifier']
        clean_metadata['id'] = self._get_first_value(self._get_first_match(metadata, id_candidates))

        # Additional enhanced fields for historical documents
        collections_candidates = ['Collections', 'COLLECTIONS', 'collections']
        clean_metadata['collections'] = self._get_first_match(metadata, collections_candidates)

        category_candidates = ['Category', 'CATEGORY', 'category']
        clean_metadata['category'] = self._get_first_match(metadata, category_candidates)

        typology_candidates = ['Typology', 'TYPOLOGY', 'typology', 'Type', 'TYPE', 'type']
        clean_metadata['typology'] = self._get_first_match(metadata, typology_candidates)

        topic_candidates = ['Topic', 'TOPIC', 'topic', 'Subject', 'SUBJECT', 'subject']
        clean_metadata['topic'] = self._get_first_match(metadata, topic_candidates)

        clean_metadata = {k: v for k, v in clean_metadata.items() if v and v != 'N/A'}

        return clean_metadata

    def display_enhanced_metadata(self, doc: Dict[str, Any], show_all_fields: bool = False):
        """
        Display metadata information for a document
        """
        metadata = doc.get("metadata", {})
        base_filename = doc["filename"].split('_chunk_')[0] if '_chunk_' in doc["filename"] else doc["filename"]

        print(f"\n=== Metadata for {base_filename} ===")

        if not metadata:
            print("No metadata available for this document.")
            return

        # Essential fields
        essential_fields = {
            'title': 'Title',
            'author': 'Author',
            'date': 'Date',
            'publisher': 'Publisher',
            'pubplace': 'Publication Place',
            'language': 'Language'
        }

        # Additional fields
        additional_fields = {
            'collections': 'Collections',
            'category': 'Category',
            'typology': 'Document Type',
            'topic': 'Topic/Subject',
            'id': 'Identifier'
        }

        # Show essential fields
        for field, label in essential_fields.items():
            if field in metadata and metadata[field]:
                print(f"{label}: {metadata[field]}")

        # Show additional fields
        if show_all_fields:
            print("\nAdditional Information:")
            for field, label in additional_fields.items():
                if field in metadata and metadata[field]:
                    print(f"{label}: {metadata[field]}")

    # Query Embeddings
    def _embed_text(self, text: str) -> List[float]:
        """Generate embeddings for input"""
        embedding = embedding_model.encode(text)
        embedding = embedding / (embedding ** 2).sum() ** 0.5
        return embedding.tolist()

    # character handling
    def _escape_query(self, query: str) -> str:
        """Escape special characters in query text for TEXT_MATCH"""
        return query.replace("'", "''")

    # User query expansion - open to suggestions
    def expand_query(self, query: str) -> str:
        """
        Use the LLM to rewrite the query for better search recall, with memory awareness
        """
        # If query expansion is disabled, return original query
        if not self.query_expansion_enabled:
            return query

        # Determine prompt based on expansion mode
        if self.query_expansion_mode == "conservative":
            system_prompt = """You are a helpful assistant that rewrites queries for better search accuracy. 
            Extract the core search intent from this query in a short phrase.
            RULES:
            - Identify only the main topic/concept/entity
            - Keep it EXTREMELY concise (5-10 words max)
            - Don't add concepts not in the original
            - Focus on nouns, proper names, and specific terms
            - If the original is already concise, return it unchanged
            - Return ONLY the search query without additional text
            - make it easy for an LLM and RAG system to understand"""

        elif self.query_expansion_mode == "moderate":
            system_prompt = """You are a helpful assistant that rewrites queries for better search accuracy.
            Extract the core search intent from this query in a moderate phrase.
            Your job is to make moderate improvements to the original query:
            RULES:
            - Identify only the main topic/concept/entity
            - Keep it concise (10-15 words max)
            - Don't add concepts not in the original
            - Focus on nouns, proper names, and specific terms
            - If the original is already concise, return it unchanged
            - Return ONLY the search query without additional text
            - make it easy for an LLM and RAG system to understand"""

        else:  # aggressive
            system_prompt = """You are a helpful assistant that rewrites queries for better search accuracy.
            Extract the core search intent from this query in a longer phrase.
            Your job is to expand the query to improve recall:
            RULES:
            - Identify only the main topic/concept/entity
            - Don't add concepts not in the original
            - Focus on nouns, proper names, and specific terms
            - make it easy for an LLM and RAG system to understand"""

        user_prompt = f"Rewrite the following query for better search recall: \n\n {query}. \n\n If appropriate, maintain references to entities mentioned in previous conversation."

        request_data = {
            "model": LLM_MODEL_NAME,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            "temperature": 0.2,
            "max_tokens": 100
        }

        try:
            response = requests.post(self.llm_api_url, json=request_data)
            if response.status_code == 200:
                expanded = response.json()["choices"][0]["message"]["content"].strip()

                # For conservative mode, if the expansion is too long, fall back to original
                if self.query_expansion_mode == "conservative" and len(expanded) > len(query) * 5:
                    print(f"Expansion too verbose for conservative mode, using original query")
                    return query

                print(f"Original query: {query}")
                print(f"Expanded query: {expanded}")
                return expanded
        except Exception as e:
            print(f"Query expansion failed: {e}")
        return query

    def _ensure_metadata_loaded(self, filenames: List[str]):
        """
        Lazy load metadata only for files that were retrieved
        """
        if not self.files_dir or not self.meta_dir:
            return

        missing_files = []
        for filename in filenames:
            base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename
            if base_filename not in self.metadata:
                missing_files.append(base_filename)

        if missing_files:
            # Remove duplicates
            missing_files = list(set(missing_files))
            print(f"Lazy loading metadata for {len(missing_files)} files...")

            for base_filename in missing_files:
                file_path = self.files_dir / base_filename
                if file_path.exists():
                    metadata = self._load_metadata(file_path)
                    if metadata:
                        self.metadata[base_filename] = metadata

            # Update cache
            self.metadata_cache[self.collection_name] = self.metadata.copy()

    # Document searching in Milvus --> (R)AG
    def retrieve(self, query: str, text_match_filter: Optional[str] = None,
                 limit: Optional[int] = None, apply_feedback: bool = False) -> List[Dict[str, Any]]:
        """
        Standard retrieve method - temporal filtering now handled separately in temporal_analyzer
        """

        if limit is None:
            limit = self.search_limit

        # Generate dense embedding for query
        query_embedding = self._embed_text(query)

        ef_value = limit * 3
        limit_multiplier = 3

        # Handle text match filter only
        expr = None
        if text_match_filter:
            escaped_filter = self._escape_query(text_match_filter)
            words = escaped_filter.split()

            if len(words) > 1:
                word_expressions = []
                for word in words:
                    word_expressions.append(f"TEXT_MATCH(text, '{word}')")

                expr = " && ".join(word_expressions)
                expr = f"({expr}) && TEXT_MATCH(text, '{escaped_filter}')"
            else:
                expr = f"TEXT_MATCH(text, '{escaped_filter}')"

        try:
            from pymilvus import AnnSearchRequest, RRFRanker, WeightedRanker

            dense_params = {
                "data": [query_embedding],
                "anns_field": "dense",
                "param": {
                    "metric_type": "COSINE",
                    "params": {"ef": ef_value}
                },
                "limit": limit * limit_multiplier
            }

            if expr:
                dense_params["expr"] = expr

            dense_request = AnnSearchRequest(**dense_params)

            sparse_params = {
                "data": [query],
                "anns_field": "sparse",
                "param": {
                    "metric_type": "BM25",
                    "params": {"drop_ratio_search": 0.2}
                },
                "limit": limit * limit_multiplier
            }

            if expr:
                sparse_params["expr"] = expr

            sparse_request = AnnSearchRequest(**sparse_params)

            ranker = WeightedRanker(self.dense_weight, self.sparse_weight)

            # Execute hybrid search
            start_time = time.time()
            results = self.milvus_client.hybrid_search(
                collection_name=self.collection_name,
                reqs=[dense_request, sparse_request],
                ranker=ranker,
                output_fields=["filename", "text"],
                limit=limit * limit_multiplier
            )
            search_time = time.time() - start_time
            print(f"Search completed in {search_time:.4f} seconds")

            processed_results = []
            seen_files = set()

            for hit in results[0]:
                filename = hit['entity']['filename']
                text = hit['entity']['text']
                score = hit.get('distance', 0.0)

                base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename
                document_metadata = self.metadata.get(base_filename, {})

                if base_filename not in seen_files and len(processed_results) < limit:
                    seen_files.add(base_filename)
                    processed_results.append({
                        "filename": filename,
                        "text": text,
                        "score": score,
                        "metadata": document_metadata
                    })

            processed_results.sort(key=lambda x: x["score"], reverse=True)

            if processed_results:
                retrieved_filenames = [result["filename"] for result in processed_results]
                self._ensure_metadata_loaded(retrieved_filenames)

                for result in processed_results:
                    base_filename = result["filename"].split('_chunk_')[0] if '_chunk_' in result["filename"] else \
                    result["filename"]
                    if base_filename in self.metadata:
                        result["metadata"] = self.metadata[base_filename]

            print(f"Retrieved {len(processed_results)} results")
            return processed_results

        except Exception as e:
            print(f"Error in hybrid search: {e}")

            # Fall back to direct text match if hybrid fails
            if text_match_filter:
                try:
                    results = self.text_match_search(text_match_filter, limit)
                    return results
                except Exception as e2:
                    print(f"Text match fallback also failed: {e2}")

            # Fall back to standard dense search
            try:
                search_params = {
                    "metric_type": "COSINE",
                    "params": {"ef": ef_value}
                }

                results = self.milvus_client.search(
                    collection_name=self.collection_name,
                    data=[query_embedding],
                    anns_field="dense",
                    search_params=search_params,
                    limit=limit,
                    output_fields=["filename", "text"],
                    expr=expr
                )

                processed_results = []
                seen_files = set()

                for hit in results[0]:
                    filename = hit["entity"]["filename"]
                    text = hit["entity"]["text"]
                    score = hit.get("distance", 0.0)

                    base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename

                    if base_filename not in seen_files and len(processed_results) < limit:
                        seen_files.add(base_filename)
                        processed_results.append({
                            "filename": filename,
                            "text": text,
                            "score": score,
                            "metadata": self.metadata.get(base_filename, {})
                        })

                print(f"Retrieved {len(processed_results)} results using fallback search")
                return processed_results

            except Exception as e3:
                print(f"All search methods failed: {e3}")
                return []

    # Entire RAG pipeline from R->A->G
    def rag_query(
            self,
            query: str,
            text_match_filter: Optional[str] = None,
            retrieval_limit: Optional[int] = None,
            use_web_search: bool = False,
            generate_report: bool = False,
            use_wos_search: bool = False,
            wos_timespan: str = "1400-01-01+1740-12-31",
    ) -> Tuple[str, Dict[str, Any]]:
        """
        Complete RAG pipeline with simplified context referencing
        """
        # Parse query for length preference
        length_preference = self._detect_length_preference(query)

        # Clean the query
        clean_query = self._clean_query(query)

        # Expand the query for better recall
        expanded_query = self.expand_query(clean_query) if self.query_expansion_enabled else clean_query

        # Retrieve relevant documents
        retrieved_docs = self.retrieve(
            query=expanded_query,
            text_match_filter=text_match_filter,
            limit=retrieval_limit
        )

        # Perform web search
        web_results = []
        if use_web_search and self.use_web_search and self.brave_api_key:
            web_results = self.web_search(clean_query, limit=2)

        if use_wos_search and WOS_API_KEY:
            wos_results = self.web_of_science_search(clean_query, limit=3, timespan=wos_timespan)

        # Combine results, prioritizing knowledge base
        combined_docs = retrieved_docs + web_results + wos_results

        # Sort by score
        combined_docs.sort(key=lambda x: x.get("score", 0), reverse=True)

        # Limit to retrieval limit
        if retrieval_limit:
            combined_docs = combined_docs[:retrieval_limit]

        if not combined_docs:
            return ("I couldn't find any relevant information to answer your question.",
                    {"error": "No relevant documents found"})

        # Check if COT should be used
        cot_used = self._should_use_cot(clean_query)

        # Generate answer using simplified method
        answer, context_data = self.generate_answer(
            clean_query, combined_docs, length_preference
        )

        # Add additional metadata to context_data
        context_data.update({
            "query": clean_query,
            "expanded_query": expanded_query if expanded_query != clean_query else None,
            "web_search_used": use_web_search and len(web_results) > 0,
            "wos_search_used": use_wos_search and len(wos_results) > 0,
            "wos_results_count": len(wos_results),
            "sources": [
                {
                    "filename": doc["filename"],
                    "score": doc.get("score", 0),
                    "metadata": doc.get("metadata", {}),
                    "source_type": "wos" if doc["filename"].startswith("wos:") else
                    "web" if doc["filename"].startswith("web:") else "knowledge_base"
                }
                for doc in combined_docs
            ]
        })

        # Generate report
        if generate_report:
            report_path, conversation_id = self.generate_report(
                query=clean_query,
                expanded_query=expanded_query,
                answer=answer,
                contexts=combined_docs,
                context_data=context_data
            )
            print(f"Report generated: {report_path}")
            context_data["report_path"] = report_path
            context_data["conversation_id"] = conversation_id

        return answer, context_data

    # Milvus 2.5 text match search implementation
    def text_match_search(self, term: str, limit: int = 20) -> List[Dict[str, Any]]:
        """
        Perform a direct text match search using TEXT_MATCH
        """
        escaped_term = self._escape_query(term)
        words = escaped_term.split()

        if len(words) > 1:
            word_expressions = []
            for word in words:
                word_expressions.append(f"TEXT_MATCH(text, '{word}')")

            expr = " && ".join(word_expressions)
            expr = f"({expr}) && TEXT_MATCH(text, '{escaped_term}')"
        else:
            # Single word query
            expr = f"TEXT_MATCH(text, '{escaped_term}')"

        print(f"Searching with expression: {expr}")

        try:
            # Use query method
            results = self.milvus_client.query(
                collection_name=self.collection_name,
                filter=expr,
                output_fields=["filename", "text"],
                limit=limit
            )

            processed_results = []
            seen_files = set()

            for doc in results:
                filename = doc["filename"]
                text = doc["text"]
                base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename

                # Only add first occurrence of each file to avoid duplicates
                if base_filename not in seen_files:
                    seen_files.add(base_filename)

                    # Get metadata for document
                    document_metadata = self.metadata.get(base_filename, {})

                    processed_results.append({
                        "filename": filename,
                        "text": text,
                        "score": 1.0,
                        "metadata": document_metadata
                    })

            # Load metadata for retrieved documents
            if processed_results:
                retrieved_filenames = [result["filename"] for result in processed_results]
                self._ensure_metadata_loaded(retrieved_filenames)

                # Update results with fresh metadata
                for result in processed_results:
                    base_filename = result["filename"].split('_chunk_')[0] if '_chunk_' in result["filename"] else \
                    result["filename"]
                    if base_filename in self.metadata:
                        result["metadata"] = self.metadata[base_filename]

            print(f"Found {len(processed_results)} results using query method")
            return processed_results

        except Exception as e:
            print(f"Query method failed: {e}")

            try:
                # Use search method with corrected parameters
                results = self.milvus_client.search(
                    collection_name=self.collection_name,
                    data=[""],
                    anns_field="sparse",
                    search_params={"metric_type": "BM25", "params": {}},
                    limit=limit,
                    expr=expr,
                    output_fields=["filename", "text"]
                )

                processed_results = []
                seen_files = set()

                for hit in results[0]:
                    filename = hit["entity"]["filename"]
                    text = hit["entity"]["text"]
                    base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename

                    if base_filename not in seen_files:
                        seen_files.add(base_filename)

                        # Get metadata for document
                        document_metadata = self.metadata.get(base_filename, {})

                        processed_results.append({
                            "filename": filename,
                            "text": text,
                            "score": hit.get("distance", 0.0),
                            "metadata": document_metadata
                        })

                print(f"Found {len(processed_results)} results using search method")
                return processed_results

            except Exception as e2:
                print(f"Search method also failed: {e2}")

                # Simple fallback without TEXT_MATCH
                try:
                    print("Attempting simple fallback search...")

                    # Use basic BM25 search on sparse field
                    results = self.milvus_client.search(
                        collection_name=self.collection_name,
                        data=[term],
                        anns_field="sparse",
                        search_params={"metric_type": "BM25"},
                        limit=limit,
                        output_fields=["filename", "text"]
                    )

                    processed_results = []
                    seen_files = set()

                    for hit in results[0]:
                        filename = hit["entity"]["filename"]
                        text = hit["entity"]["text"]

                        # Only include if text actually contains the term (manual filter)
                        if term.lower() in text.lower():
                            base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename

                            if base_filename not in seen_files:
                                seen_files.add(base_filename)

                                document_metadata = self.metadata.get(base_filename, {})

                                processed_results.append({
                                    "filename": filename,
                                    "text": text,
                                    "score": hit.get("distance", 0.0),
                                    "metadata": document_metadata
                                })

                    print(f"Found {len(processed_results)} results using fallback method")
                    return processed_results

                except Exception as e3:
                    print(f"All search methods failed: {e3}")
                    return []

    # Brave API for web searches. --> soon to change to WoS
    def web_search(self, query: str, limit: int = 3) -> List[Dict[str, Any]]:
        """
        Perform a web search using Brave Search API
        """
        if not self.brave_api_key:
            print("Brave API key not set. Skipping web search.")
            return []

        headers = {
            "Accept": "application/json",
            "X-Subscription-Token": self.brave_api_key
        }

        params = {
            "q": query,
            "count": limit
        }

        try:
            response = requests.get(
                BRAVE_SEARCH_API_URL,
                headers=headers,
                params=params
            )

            if response.status_code == 200:
                results = response.json()

                web_results = []
                for web in results.get("web", {}).get("results", []):
                    web_results.append({
                        "filename": f"web:{web.get('url', '')}",
                        "text": f"{web.get('title', '')}\n{web.get('description', '')}\n{web.get('url', '')}",
                        "score": 0.9,
                        "doc_id": f"web:{hashlib.md5(web.get('url', '').encode()).hexdigest()[:8]}"
                    })

                return web_results
            else:
                print(f"Web search failed with status code: {response.status_code}")
                return []

        except Exception as e:
            print(f"Error performing web search: {e}")
            return []

    # Implement chain of thought for LLM
    def _should_use_cot(self, query: str) -> bool:
        """
        Determine if Chain of Thought should be enabled
        """
        # Simple heuristics to detect complex queries
        complex_indicators = [
            "why", "how", "explain", "analyze", "compare", "contrast",
            "difference", "relationship", "impact", "influence", "cause",
            "effect", "result", "consequence", "significance", "importance"
        ]

        # Check for complex indicators
        query_lower = query.lower()
        has_complex_terms = any(indicator in query_lower for indicator in complex_indicators)

        # Check for length
        is_long_query = len(query.split()) > 10

        # Check for multiple questions
        has_multiple_questions = query.count("?") > 1

        return has_complex_terms or is_long_query or has_multiple_questions

    # CRUCIAL - Generate LLM answer based on retrieval R(AG) with citations
    def generate_answer(
            self,
            query: str,
            contexts: List[Dict[str, Any]],
            length_preference: str = "normal",
            conversation_memory: Optional[List[Tuple[str, str]]] = None
    ) -> Tuple[str, Dict[str, Any]]:
        """
        Generate an answer with influence analysis detection
        """

        # Create simple context references
        context_refs = {}
        context_texts = []

        for i, doc in enumerate(contexts):
            context_key = f"context_{i + 1}"
            context_refs[context_key] = {
                "filename": doc["filename"],
                "text": doc["text"][:500] + ("..." if len(doc["text"]) > 500 else ""),
                "score": doc.get("score", 0.0),
                "metadata": doc.get("metadata", {})
            }

            # Add numbered context for LLM
            context_text = f"[Context {i + 1}] {doc['text']}"
            context_texts.append(context_text)

        # Prepare context for LLM (limit length based on preference)
        max_context_tokens = 1500 if length_preference == "short" else 4000 if length_preference == "detailed" else 2500
        snippet_size = 1000 if length_preference == "short" else 2500 if length_preference == "detailed" else 1500

        snippets = []
        current_tokens = 0

        for i, doc in enumerate(contexts):
            # Get the full text from the document
            doc_text = doc["text"]
            context_label = f"[Context {i + 1} - {doc['filename']}]"

            # Add date information to context if available
            if doc.get('metadata', {}).get('date'):
                context_label += f" (Date: {doc['metadata']['date']})"

            estimated_tokens = len(doc_text) // 4

            # If document would put us over the limit, use a snippet
            if current_tokens + estimated_tokens > max_context_tokens:
                snippet = doc_text[:snippet_size]
                snippet_text = f"{context_label}:\n{snippet}..."
                snippets.append(snippet_text)
                current_tokens += len(snippet_text) // 4
            else:
                # Use full document text
                full_text = f"{context_label}:\n{doc_text}"
                snippets.append(full_text)
                current_tokens += estimated_tokens

            # If we're at token limit, stop adding documents
            if current_tokens >= max_context_tokens:
                break

        combined_context = "\n\n".join(snippets)

        # Include conversation memory if enabled
        conversation_context = ""
        if self.memory_enabled and conversation_memory:
            recent_memory = conversation_memory[-self.memory_turns:]
            memory_text = "\n\n".join([f"User: {q}\nAssistant: {a}" for q, a in recent_memory])
            conversation_context = f"\nRecent conversation history:\n{memory_text}\n\n"

        # Check if we should use Chain of Thought
        use_cot = self._should_use_cot(query)

        system_prompt = f"""You are a research historian with expertise in 15th–17th century Europe.
            Analyze and synthesize the provided sources to answer the user's question comprehensively.
            Draw connections across sources, identify patterns, and provide historical interpretation.
            Structure your response as a cohesive argument supported by evidence rather than going through sources sequentially.
            Offer scholarly insights and contextualize findings within the broader historical period.
    
            Do NOT reference the titles of the contexts in any manner. Speak holistically as if you are speaking using
            internal knowledge."""

        if use_cot:
            system_prompt += """
            For this complex question, think step by step:
            1. Break down what information is needed
            2. Identify relevant facts from the context
            3. Consider how these facts relate to each other
            4. Form a clear conclusion
            """

        # Adjust response length
        if length_preference == "short":
            system_prompt += "\nProvide a brief and concise answer."
        elif length_preference == "detailed":
            system_prompt += "\nProvide a thorough and comprehensive answer."

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user",
             "content": f"{conversation_context}Here is the context information:\n\n{combined_context}\n\nQuestion: {query}"}
        ]

        # Adjust max_tokens based on preferences
        adjusted_max_tokens = 256 if length_preference == "short" else self.max_tokens * 2 if length_preference == "detailed" else self.max_tokens
        if use_cot:
            adjusted_max_tokens = int(adjusted_max_tokens * 1.5)

        request_data = {
            "model": LLM_MODEL_NAME,
            "messages": messages,
            "temperature": self.temperature,
            "max_tokens": adjusted_max_tokens
        }

        print(f"[LLM DEBUG] Request parameters:")
        print(f"  Temperature: {self.temperature}")
        print(f"  Max tokens: {adjusted_max_tokens}")
        print(f"  Context length: {len(combined_context)} chars")
        print(f"  Number of contexts: {len(snippets)}")
        print(f"  CoT enabled: {use_cot}")

        try:
            response = requests.post(self.llm_api_url, json=request_data)

            if response.status_code == 200:
                result = response.json()
                answer = result["choices"][0]["message"]["content"]

                context_data = {
                    "contexts": context_refs,
                    "cot_used": use_cot,
                    "context_count": len(contexts),
                    "tokens_used": adjusted_max_tokens
                }
                return answer, context_data
            else:
                print(f"LLM API error: {response.status_code}")
                print(response.text)
                error_msg = f"Error: Failed to generate response (Status code: {response.status_code})"
                return error_msg, {"error": error_msg}

        except Exception as e:
            print(f"Exception when calling LLM API: {e}")
            error_msg = f"Error: Failed to generate response due to an exception: {e}"
            return error_msg, {"error": error_msg}

    # Understand which documents were used in generation
    def generate_context_visualization(self, contexts: List[Dict[str, Any]], citation_data: Dict) -> Dict:
        """
        Generate visualization data showing which parts of documents were used
        """
        if not contexts or not citation_data:
            return {"error": "No context or citation data available"}

        # Create visualization data
        viz_data = {
            "documents": [],
            "sections": citation_data.get("sections", []),
            "heatmap": {}
        }

        # Process each context
        for ctx_id, ctx_info in citation_data.get("contexts", {}).items():
            doc_data = {
                "id": ctx_id,
                "filename": ctx_info["filename"],
                "preview": ctx_info["text"][:150] + "...",
                "relevance_score": ctx_info.get("score", 0)
            }
            viz_data["documents"].append(doc_data)

            # Track citation frequency for heatmap
            viz_data["heatmap"][ctx_id] = 0

        # Count citations for heatmap
        for section in citation_data.get("sections", []):
            for ctx_id in section.get("citations", []):
                if ctx_id in viz_data["heatmap"]:
                    viz_data["heatmap"][ctx_id] += 1

        return viz_data

    # Change length of answer based on keywords
    def _detect_length_preference(self, query: str) -> str:
        """
        Detect if the user explicitly wants a short or detailed response
        """
        query_lower = query.lower()

        # Check for indicators of short response
        short_indicators = [
            "make it short", "brief", "concise", "summarize", "short version",
            "quick answer", "tldr", "summary", "briefly", "in short", "shorter", "quickly",
            "keep it simple", "in a few words", "brief overview", "short summary", "shortly", "concisely"
        ]

        # Check for indicators of detailed response
        detailed_indicators = [
            "detailed", "comprehensive", "thorough", "in-depth", "elaborate",
            "exhaustive", "explain fully", "complete answer", "full explanation",
            "make it longer", "longer", "make it long", "explain everything",
            "explain in detail", "be specific", "make it very long"
        ]

        for indicator in short_indicators:
            if indicator in query_lower:
                return "short"

        for indicator in detailed_indicators:
            if indicator in query_lower:
                return "detailed"

        return "normal"

    # Remove identifiers for LLM input
    def _clean_query(self, query: str) -> str:
        """
        Remove length modifiers from the query
        """
        # List of phrases to remove
        phrases_to_remove = [
            "make it short", "brief", "concise", "summarize", "short version",
            "quick answer", "tldr", "summary", "briefly", "in short",
            "detailed", "comprehensive", "thorough", "in-depth", "elaborate",
            "exhaustive", "explain fully", "complete answer", "full explanation",
            "make it longer", "longer", "make it long", "shorter", "quickly",
            "keep it simple", "in a few words", "brief overview", "short summary",
            "explain everything", "explain in detail", "be specific", "make it very long",
            "shortly", "concisely"
        ]

        cleaned_query = query
        for phrase in phrases_to_remove:
            # Replace phrases with empty string
            cleaned_query = re.sub(r'\b' + re.escape(phrase) + r'\b', '', cleaned_query, flags=re.IGNORECASE)

        # Clean up double spaces
        cleaned_query = re.sub(r'\s+', ' ', cleaned_query).strip()

        return cleaned_query

    def display_retrieved_contexts(self, contexts: List[Dict[str, Any]], highlight_term: Optional[str] = None):
        """
        Display retrieved contexts with optional highlighting
        """
        if not contexts:
            print("No contexts retrieved.")
            return

        for i, doc in enumerate(contexts):
            print(f"\n--- Context {i + 1} ---")
            print(f"Source: {doc['filename']}")
            if "score" in doc:
                print(f"Score: {doc['score']:.4f}")
                if "original_score" in doc and doc["score"] != doc["original_score"]:
                    print(f"Original Score: {doc['original_score']:.4f} (adjusted by feedback)")

            text = doc["text"]

            # Find occurrences
            if highlight_term:
                term_positions = []
                start_pos = 0
                text_lower = text.lower()
                highlight_term_lower = highlight_term.lower()

                while True:
                    pos = text_lower.find(highlight_term_lower, start_pos)
                    if pos == -1:
                        break
                    term_positions.append(pos)
                    start_pos = pos + len(highlight_term)

                if term_positions:
                    print(f"Contains {len(term_positions)} occurrences of '{highlight_term}'")

                    # Show context for first occurrence
                    pos = term_positions[0]
                    start = max(0, pos - 100)
                    end = min(len(text), pos + len(highlight_term) + 100)

                    context = text[start:end]
                    # Highlight the term
                    term_text = text[pos:pos + len(highlight_term)]
                    highlighted = context.replace(term_text, f"**{term_text}**", 1)

                    print(f"Preview: ...{highlighted}...")
                else:
                    text_preview = text[:200].replace('\n', ' ').strip()
                    print(f"Preview: {text_preview}...")
            else:
                text_preview = text[:200].replace('\n', ' ').strip()
                print(f"Preview: {text_preview}...")

            # Show total text length
            print(f"Text length: {len(text)} characters")

            # Show citation ID for reference
            if "doc_id" in doc:
                print(f"Document ID: {doc['doc_id']}")

    # File IO for Researchers to view txt file of entire conversation
    def generate_report(
            self,
            query: str,
            expanded_query: str,
            answer: str,
            contexts: List[Dict[str, Any]],
            context_data: Dict,
            conversation_id: str = None
    ) -> Tuple[str, str]:
        """
        Generate simplified text report of RAG query and results
        """
        # Create report directory
        report_dir = "rag_reports"
        os.makedirs(report_dir, exist_ok=True)

        # Create/use conversation ID
        if conversation_id is None:
            conversation_id = f"conversation_{time.strftime('%Y%m%d-%H%M%S')}"

        # Define filepath
        filepath = os.path.join(report_dir, f"{conversation_id}.txt")
        is_new_report = not os.path.exists(filepath)

        with open(filepath, "a", encoding="utf-8") as f:
            # Write header for new reports
            if is_new_report:
                f.write("=" * 80 + "\n")
                f.write(f"RAG SYSTEM CONVERSATION REPORT - Started {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"Collection: {self.collection_name}\n")
                f.write("=" * 80 + "\n\n")

            # Add timestamp for each query
            f.write("\n" + "=" * 80 + "\n")
            query_number = sum(1 for line in open(filepath, 'r') if 'QUERY:' in line) + 1 if not is_new_report else 1
            f.write(f"QUERY #{query_number} - {time.strftime('%H:%M:%S')}\n")
            f.write("=" * 80 + "\n\n")

            # Query information
            f.write("QUERY: " + query + "\n")
            if expanded_query and query != expanded_query:
                f.write("EXPANDED QUERY: " + expanded_query + "\n")
            f.write("\n")

            # System settings
            f.write("SETTINGS: ")
            settings = []
            if context_data.get('memory_enabled', False):
                settings.append("Memory")
            if context_data.get('query_expansion_enabled', False):
                settings.append("Query Expansion")
            if context_data.get('web_search_used', False):
                settings.append("Web Search")
            if context_data.get('cot_used', False):
                settings.append("Chain of Thought")
            f.write(", ".join(settings) + "\n\n")

            # Answer
            f.write("ANSWER:\n")
            f.write("-" * 80 + "\n")
            f.write(f"{answer}\n\n")

            # Context sources
            f.write("SOURCES USED:\n")
            f.write("-" * 80 + "\n")

            for i, context_info in enumerate(context_data.get('contexts', {}).values()):
                f.write(f"Context {i + 1}: {context_info['filename']}")

                metadata = context_info.get("metadata", {})
                if metadata:
                    meta_parts = []
                    if metadata.get("title"):
                        meta_parts.append(metadata['title'])
                    if metadata.get("author"):
                        meta_parts.append(f"by {metadata['author']}")
                    if metadata.get("date"):
                        meta_parts.append(f"({metadata['date']})")

                    if meta_parts:
                        f.write(f" - {' '.join(meta_parts)}")

                f.write(f" (Score: {context_info.get('score', 0):.2f})\n")

            # Query statistics
            f.write(f"\nSTATISTICS:\n")
            f.write(f"- Contexts retrieved: {context_data.get('context_count', 0)}\n")
            f.write(f"- CoT reasoning: {'Yes' if context_data.get('cot_used', False) else 'No'}\n")
            f.write(f"- Collection: {self.collection_name}\n")
            f.write(f"- Length preference: {context_data.get('length_preference', 'normal')}\n")

            f.write("\n")

        return filepath, conversation_id

    def generate_streaming_answer(
            self,
            query: str,
            contexts: List[Dict[str, Any]],
            length_preference: str = "normal",
            conversation_memory: Optional[List[Tuple[str, str]]] = None
    ) -> Iterator[Dict[str, Any]]:
        """
        Generate a streaming answer
        """
        # Create simple context references
        context_refs = {}
        context_texts = []

        for i, doc in enumerate(contexts):
            context_key = f"context_{i + 1}"
            context_refs[context_key] = {
                "filename": doc["filename"],
                "text": doc["text"][:500] + ("..." if len(doc["text"]) > 500 else ""),
                "score": doc.get("score", 0.0),
                "metadata": doc.get("metadata", {})
            }

            # Add numbered context for LLM
            context_text = f"[Context {i + 1}] {doc['text']}"
            context_texts.append(context_text)

        # Prepare context for LLM (limit length based on preference)
        max_context_tokens = 1500 if length_preference == "short" else 4000 if length_preference == "detailed" else 2500
        snippet_size = 1000 if length_preference == "short" else 2500 if length_preference == "detailed" else 1500

        snippets = []
        current_tokens = 0

        for i, doc in enumerate(contexts):
            # Get the full text from the document
            doc_text = doc["text"]
            context_label = f"[Context {i + 1} - {doc['filename']}]"

            estimated_tokens = len(doc_text) // 4

            # If document would put us over the limit, use a snippet
            if current_tokens + estimated_tokens > max_context_tokens:
                snippet = doc_text[:snippet_size]
                snippet_text = f"{context_label}:\n{snippet}..."
                snippets.append(snippet_text)
                current_tokens += len(snippet_text) // 4
            else:
                # Use full document text
                full_text = f"{context_label}:\n{doc_text}"
                snippets.append(full_text)
                current_tokens += estimated_tokens

            # If we're at token limit, stop adding documents
            if current_tokens >= max_context_tokens:
                break

        combined_context = "\n\n".join(snippets)

        # Include conversation memory if enabled
        conversation_context = ""
        if self.memory_enabled and conversation_memory:
            recent_memory = conversation_memory[-self.memory_turns:]
            memory_text = "\n\n".join([f"User: {q}\nAssistant: {a}" for q, a in recent_memory])
            conversation_context = f"\nRecent conversation history:\n{memory_text}\n\n"
        use_cot = self._should_use_cot(query)

        # Simplified system prompt
        system_prompt = f"""You are a research historian with expertise in 15th–17th century Europe.
        Analyze and synthesize the provided sources to answer the user's question comprehensively.
        Draw connections across sources, identify patterns, and provide historical interpretation.
        Structure your response as a cohesive argument supported by evidence rather than going through sources sequentially.
        Offer scholarly insights and contextualize findings within the broader historical period.

        Do NOT reference the titles of the contexts in any manner. Speak holistically as if you are speaking using
        internal knowledge."""

        if use_cot:
            system_prompt += """
            For this complex question, think step by step:
            1. Break down what information is needed
            2. Identify relevant facts from the context
            3. Consider how these facts relate to each other
            4. Form a clear conclusion
            """

        # Adjust response length
        if length_preference == "short":
            system_prompt += "\nProvide a brief and concise answer."
        elif length_preference == "detailed":
            system_prompt += "\nProvide a thorough and comprehensive answer."

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user",
             "content": f"{conversation_context}Here is the context information:\n\n{combined_context}\n\nQuestion: {query}"}
        ]

        # Adjust max_tokens based on preferences
        adjusted_max_tokens = 256 if length_preference == "short" else self.max_tokens * 2 if length_preference == "detailed" else self.max_tokens
        if use_cot:
            adjusted_max_tokens = int(adjusted_max_tokens * 1.5)

        request_data = {
            "model": LLM_MODEL_NAME,
            "messages": messages,
            "temperature": self.temperature,
            "max_tokens": adjusted_max_tokens,
            "stream": True
        }

        print(f"[LLM DEBUG] Streaming request parameters:")
        print(f"  Temperature: {self.temperature}")
        print(f"  Max tokens: {adjusted_max_tokens}")
        print(f"  Context length: {len(combined_context)} chars")
        print(f"  Number of contexts: {len(snippets)}")
        print(f"  CoT enabled: {use_cot}")

        try:
            import requests
            response = requests.post(
                self.llm_api_url,
                json=request_data,
                stream=True,
                timeout=60
            )

            if response.status_code == 200:
                full_answer = ""

                # First, send context information
                yield {
                    "type": "context",
                    "contexts": context_refs,
                    "cot_used": use_cot,
                    "context_count": len(contexts)
                }

                # Stream the response
                for line in response.iter_lines():
                    if line:
                        line = line.decode('utf-8')
                        if line.startswith('data: '):
                            data = line[6:]
                            if data.strip() == '[DONE]':
                                break

                            try:
                                chunk_data = json.loads(data)
                                if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
                                    choice = chunk_data['choices'][0]
                                    if 'delta' in choice and 'content' in choice['delta']:
                                        content = choice['delta']['content']
                                        if content:
                                            full_answer += content
                                            yield {
                                                "type": "content",
                                                "content": content
                                            }
                            except json.JSONDecodeError:
                                continue

                # Send final completion message
                yield {
                    "type": "done",
                    "full_answer": full_answer,
                    "context_data": {
                        "contexts": context_refs,
                        "cot_used": use_cot,
                        "context_count": len(contexts),
                        "tokens_used": adjusted_max_tokens,
                        "query": query,
                        "length_preference": length_preference,
                        "sources": [
                            {
                                "filename": doc["filename"],
                                "score": doc.get("score", 0),
                                "metadata": doc.get("metadata", {})
                            }
                            for doc in contexts
                        ]
                    }
                }

            else:
                print(f"LLM API error: {response.status_code}")
                print(response.text)
                error_msg = f"Error: Failed to generate response (Status code: {response.status_code})"
                yield {
                    "type": "error",
                    "error": error_msg
                }

        except Exception as e:
            print(f"Exception when calling LLM API: {e}")
            error_msg = f"Error: Failed to generate response due to an exception: {e}"
            yield {
                "type": "error",
                "error": error_msg
            }

    # Document clustering code
    def find_similar_documents(self, filename: str, limit: int = 5) -> List[Dict[str, Any]]:
        """Find documents similar to the given document"""
        try:
            # Get document embedding
            results = self.milvus_client.query(
                collection_name=self.collection_name,
                filter=f"filename == '{filename.replace(chr(39), chr(39) + chr(39))}'",
                output_fields=["dense"],
                limit=1
            )

            if not results:
                return []

            source_embedding = results[0]["dense"]

            # Find similar documents
            similar_results = self.milvus_client.search(
                collection_name=self.collection_name,
                data=[source_embedding],
                anns_field="dense",
                search_params={"metric_type": "COSINE", "params": {"ef": 50}},
                limit=limit + 10,
                output_fields=["filename", "text"]
            )

            # Process results
            similar_docs = []
            seen_files = set()
            source_base = filename.split('_chunk_')[0] if '_chunk_' in filename else filename

            for hit in similar_results[0]:
                hit_filename = hit["entity"]["filename"]
                base_filename = hit_filename.split('_chunk_')[0] if '_chunk_' in hit_filename else hit_filename

                if (base_filename != source_base and
                        base_filename not in seen_files and
                        len(similar_docs) < limit):
                    seen_files.add(base_filename)

                    similar_docs.append({
                        "filename": hit_filename,
                        "text": hit["entity"]["text"],
                        "score": hit["distance"],
                        "metadata": self.metadata.get(base_filename, {})
                    })

            return similar_docs

        except Exception as e:
            print(f"Error finding similar documents: {e}")
            return []

    def get_document_clusters(self, n_clusters: int = 5, sample_size: int = 100) -> Dict:
        """Cluster documents by similarity"""
        try:
            # Get sample documents
            results = self.milvus_client.query(
                collection_name=self.collection_name,
                filter="",
                output_fields=["filename", "dense", "text"],
                limit=sample_size
            )

            if len(results) < n_clusters:
                return {"error": "Not enough documents for clustering"}

            from sklearn.cluster import KMeans
            import numpy as np

            # Extract data
            embeddings = np.array([doc["dense"] for doc in results])

            # Perform clustering
            kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
            cluster_labels = kmeans.fit_predict(embeddings)

            # Group results
            clusters = {}
            for i, label in enumerate(cluster_labels):
                if label not in clusters:
                    clusters[label] = []

                doc = results[i]
                base_filename = doc["filename"].split('_chunk_')[0] if '_chunk_' in doc["filename"] else doc["filename"]

                clusters[label].append({
                    "filename": doc["filename"],
                    "text_preview": doc["text"][:200] + "...",
                    "metadata": self.metadata.get(base_filename, {})
                })

            return {"clusters": clusters, "total_documents": len(results)}

        except ImportError:
            return {"error": "scikit-learn required: pip install scikit-learn"}
        except Exception as e:
            return {"error": str(e)}

    # Temporal Analysis
    def temporal_rag_query(self, query: str, date_range: Optional[Tuple[int, int]] = None,
                           **kwargs) -> Tuple[str, Dict[str, Any]]:
        """RAG query with temporal analysis capabilities"""

        # Parse temporal information
        if not date_range:
            temporal_info = self.temporal_analyzer.parse_date_query(query)
            if temporal_info['date_range']:
                date_range = temporal_info['date_range']
            elif temporal_info['specific_year']:
                year = temporal_info['specific_year']
                date_range = (year, year)
            elif temporal_info['decade']:
                date_range = temporal_info['decade']
            elif temporal_info['century']:
                date_range = temporal_info['century']

        # Clean query
        length_preference = self._detect_length_preference(query)
        clean_query = self._clean_query(query)
        expanded_query = self.expand_query(clean_query) if self.query_expansion_enabled else clean_query

        # Retrieve with temporal constraints
        if date_range:
            start_year, end_year = date_range
            retrieved_docs = self.temporal_analyzer.retrieve_with_temporal_filter(
                query=expanded_query,
                start_year=start_year,
                end_year=end_year,
                limit=kwargs.get('retrieval_limit', self.search_limit)
            )
        else:
            retrieved_docs = self.retrieve(
                query=expanded_query,
                limit=kwargs.get('retrieval_limit', self.search_limit)
            )

        # Web search
        web_results = []
        if kwargs.get('use_web_search', False) and self.use_web_search and self.brave_api_key:
            web_results = self.web_search(clean_query, limit=2)

        combined_docs = retrieved_docs + web_results
        combined_docs.sort(key=lambda x: x.get("score", 0), reverse=True)

        if kwargs.get('retrieval_limit'):
            combined_docs = combined_docs[:kwargs['retrieval_limit']]

        if not combined_docs:
            return ("I couldn't find any relevant information for that time period.",
                    {"error": "No relevant documents found", "date_range": date_range})

        # Generate temporal aware answer
        answer, context_data = self.generate_answer(clean_query, combined_docs, length_preference)

        # Add temporal context
        context_data.update({
            "query": clean_query,
            "expanded_query": expanded_query if expanded_query != clean_query else None,
            "date_range": date_range,
            "temporal_analysis": True
        })

        return answer, context_data

    def _check_temporal_match(self, date_str: str, temporal_filter: str) -> bool:
        """Check if date matches temporal filter"""
        if not date_str or not temporal_filter:
            return True

        try:
            year = int(date_str)
            # Simple temporal filter
            if ">=" in temporal_filter and "<=" in temporal_filter:
                # Extract date range
                parts = temporal_filter.split(" and ")
                start_part = parts[0]
                end_part = parts[1]

                start_year = int(start_part.split("'")[1])
                end_year = int(end_part.split("'")[1])

                return start_year <= year <= end_year
            elif ">=" in temporal_filter:
                start_year = int(temporal_filter.split("'")[1])
                return year >= start_year
            elif "<=" in temporal_filter:
                end_year = int(temporal_filter.split("'")[1])
                return year <= end_year
            elif "==" in temporal_filter:
                target_year = int(temporal_filter.split("'")[1])
                return year == target_year

        except (ValueError, IndexError):
            pass

        return True

    # Web of Science API addition
    def web_of_science_search(self, query: str, limit: int = 5, timespan: str = "1400-01-01+1700-12-31") -> List[Dict[str, Any]]:
        if not self.wos_api_key:
            print("WOS API Key not available")
            return []

        headers = {
            "X-ApiKey": self.wos_api_key,
            "Accept": "application/json"
        }

        # Use the documents endpoint
        search_url = f"{self.wos_api_url}/documents"

        # Convert timespan to year range for WoS query
        try:
            year_range = timespan.replace("-01-01+", "-").replace("-12-31", "")
            start_year, end_year = year_range.split("-")

            # Build proper WoS query with field tags
            if int(end_year) < 1900:
                # For historical periods, search without date filter since WoS has few historical papers
                wos_query = f"TS=({query})"
                print(f"Historical period detected ({start_year}-{end_year}), searching without date filter")
            else:
                # For modern periods, include date range
                wos_query = f"TS=({query}) AND PY=({start_year}-{end_year})"

        except:
            # Fallback if timespan parsing fails
            wos_query = f"TS=({query})"

        params = {
            "q": wos_query,
            "limit": min(limit, 50),
            "page": 1
        }

        try:
            print(f"Searching Web of Science at: {search_url}")
            print(f"Query: {wos_query}")

            response = requests.get(
                search_url,
                headers=headers,
                params=params,
                timeout=30
            )

            print(f"WoS API Status: {response.status_code}")

            if response.status_code == 200:
                data = response.json()

                wos_results = []
                hits = data.get("hits", [])
                metadata = data.get("metadata", {})

                print(f"Found {len(hits)} WoS results out of {metadata.get('total', 'unknown')} total")

                for i, hit in enumerate(hits):
                    try:
                        source = hit.get("source", {})

                        # Extract title
                        title = hit.get("title", "")
                        if isinstance(title, list) and title:
                            title = title[0]

                        if not title:
                            print(f"Skipping result {i + 1} - no title")
                            continue

                        print(f"Processing result {i + 1}: {title[:80]}...")

                        # Extract authors from 'names' field
                        authors = []
                        names = hit.get("names", {})
                        if "authors" in names:
                            for author in names["authors"]:
                                if isinstance(author, dict):
                                    # Try different author name fields
                                    name = author.get("displayName", author.get("name", author.get("fullName", "")))
                                    if name:
                                        authors.append(name)
                                elif isinstance(author, str):
                                    authors.append(author)

                        # Extract journal/source from source object
                        journal = source.get("sourceTitle", "")

                        # Extract publication year from source
                        pub_year = str(source.get("publishYear", ""))

                        # Extract DOI from identifiers
                        doi = ""
                        identifiers = hit.get("identifiers", {})
                        if "doi" in identifiers:
                            doi = identifiers["doi"]
                        elif isinstance(identifiers, list):
                            for identifier in identifiers:
                                if isinstance(identifier, dict) and identifier.get("type") == "doi":
                                    doi = identifier.get("value", "")
                                    break

                        # Extract abstract - might be in different places
                        abstract = hit.get("abstract", "")
                        if not abstract and "abstracts" in hit:
                            abstracts = hit["abstracts"]
                            if isinstance(abstracts, list) and abstracts:
                                abstract = abstracts[0]

                        # Extract citation count
                        citations_info = hit.get("citations", {})
                        times_cited = 0
                        if isinstance(citations_info, dict):
                            times_cited = citations_info.get("count", 0)
                        elif isinstance(citations_info, int):
                            times_cited = citations_info

                        # Get UID
                        uid = hit.get("uid", hashlib.md5(title.encode()).hexdigest()[:8])

                        formatted_result = {
                            "filename": f"wos:{uid}",
                            "text": self._format_wos_document_text(title, authors, journal, pub_year, abstract, doi),
                            "score": 0.95,
                            "metadata": {
                                "title": title,
                                "authors": authors,
                                "journal": journal,
                                "publication_year": pub_year,
                                "doi": doi,
                                "source": "Web of Science",
                                "uid": uid,
                                "times_cited": times_cited
                            },
                            "doc_id": f"wos:{uid}"
                        }
                        wos_results.append(formatted_result)

                        print(
                            f"Successfully processed: {title[:50]}... ({len(authors)} authors, {journal}, {pub_year})")

                    except Exception as e:
                        print(f"Error processing WoS record {i + 1}: {e}")
                        continue

                print(f"Successfully processed {len(wos_results)} WoS results")
                return wos_results

            elif response.status_code == 401:
                print("Web of Science API authentication failed.")
                return []
            elif response.status_code == 400:
                print("Web of Science API bad request.")
                print(f"Query used: {wos_query}")
                print(f"Response: {response.text}")
                return []
            else:
                print(f"Web of Science API error: {response.status_code}")
                print(f"Response: {response.text}")
                return []

        except Exception as e:
            print(f"Error performing Web of Science search: {e}")
            return []

    def _format_wos_document_text(self, title: str, authors: List[str], journal: str, year: str, abstract: str, doi: str) -> str:
        """
        Format Web of Science document
        """
        text_parts = []

        if title:
            text_parts.append(f"Title: {title}")

        if authors:
            text_parts.append(f"Authors: {', '.join(authors)}")

        if journal:
            text_parts.append(f"Journal: {journal}")

        if year:
            text_parts.append(f"Publication Year: {year}")

        if doi:
            text_parts.append(f"DOI: {doi}")

        if abstract:
            text_parts.append(f"Abstract: {abstract}")

        return "\n\n".join(text_parts)

# Deprecated since version 1.7
def main():
    """Interactive RAG CLI application"""
    print("Initializing RAG system v1.9...")
    try:
        rag = RAGSystem(
            memory_enabled=True,
            memory_turns=3,
            query_expansion_enabled=True
        )
        print("RAG system initialized successfully!")
    except Exception as e:
        print(f"Failed to initialize RAG system: {e}")
        return

    print("\n" + "=" * 50)
    print("Welcome to the Enhanced RAG System v1.9")
    print("- Conversation memory and chain-of-thought reasoning")
    print("- Interactive citations and context visualization")
    print("- Query expansion (now toggleable)")
    print("- Web search integration and feedback system")
    print("=" * 50)

    while True:
        print("\nOptions:")
        print("1. Ask a question (RAG query with citations)")
        print("2. Retrieve documents only")
        print("3. Advanced search with text match filter")
        print("4. Toggle web search (currently", "ON" if rag.use_web_search else "OFF", ")")
        print("5. Memory settings")
        print("6. Toggle query expansion (currently", "ON" if rag.query_expansion_enabled else "OFF", ")")
        print("7. Configure search parameters")
        print("8. Exit")

        choice = input("\nEnter your choice (1-8): ").strip()

        if choice == "1":
            query = input("\nEnter your question: ").strip()
            if not query:
                print("Please enter a valid question.")
                continue

            use_web = input("Include web search? (y/n, default: n): ").lower().strip() == 'y'

            print("\nProcessing your question...")
            answer, citation_data = rag.rag_query(query, use_web_search=use_web)

            print("\n" + "=" * 50)
            print("ANSWER:")
            print(answer)
            print("=" * 50)

            # Display generation ID for citation reference
            generation_id = citation_data.get("generation_id", "unknown")
            print(f"\nGeneration ID: {generation_id}")

            # Option to see citations
            show_citations = input("\nDo you want to see citation details? (y/n): ").lower().strip()
            if show_citations == "y":
                print("\nCITATION DETAILS:")
                sections = citation_data.get("sections", [])

                if not sections:
                    print("No citation data available.")
                else:
                    for section in sections:
                        print(f"\nSection {section['id']}:")
                        print(f"Text: {section['text'][:100]}...")
                        if section["citations"]:
                            print(f"Cited sources: {', '.join(section['citations'])}")

                            # Show details for first citation
                            first_citation = section["citations"][0]
                            if first_citation in citation_data.get("contexts", {}):
                                ctx = citation_data["contexts"][first_citation]
                                print(f"From: {ctx['filename']}")
                                print(f"Content: {ctx['text'][:150]}...")
                        else:
                            print("No citations for this section")


        elif choice == "2":
            query = input("\nEnter search query: ").strip()
            if not query:
                print("Please enter a valid search query.")
                continue

            limit = input("Enter number of results to retrieve (default: 5): ").strip()
            try:
                limit = int(limit) if limit else 5
            except ValueError:
                print("Invalid number, using default: 5")
                limit = 5

            print("\nRetrieving documents...")
            contexts = rag.retrieve(query, limit=limit)

            print(f"\nRetrieved {len(contexts)} documents.")
            highlight = input("Enter term to highlight in contexts (optional): ").strip()
            highlight_term = highlight if highlight else None
            rag.display_retrieved_contexts(contexts, highlight_term)

        elif choice == "3":
            query = input("\nEnter semantic search query: ").strip()
            text_filter = input("Enter text to match exactly: ").strip()

            if not query and not text_filter:
                print("Please enter at least one search criterion.")
                continue

            limit = input("Enter number of results to retrieve (default: 5): ").strip()
            try:
                limit = int(limit) if limit else 5
            except ValueError:
                print("Invalid number, using default: 5")
                limit = 5

            print("\nRetrieving documents...")
            contexts = rag.retrieve(
                query=query if query else "",
                text_match_filter=text_filter if text_filter else None,
                limit=limit
            )

            print(f"\nRetrieved {len(contexts)} documents.")
            highlight = input(f"Enter term to highlight in contexts (default: '{text_filter}'): ").strip()
            highlight_term = highlight if highlight else (text_filter if text_filter else None)
            rag.display_retrieved_contexts(contexts, highlight_term)

            # Option to generate answer
            if contexts:
                generate = input("\nGenerate an answer based on these results? (y/n): ").lower().strip()
                if generate == "y":
                    if not query:
                        query = input("Enter a question for the generation: ").strip()

                    if query:
                        print("\nGenerating answer...")
                        answer = rag.generate_answer(query, contexts)

                        print("\n" + "=" * 50)
                        print("ANSWER:")
                        print(answer)
                        print("=" * 50)

                    else:
                        print("No query provided for generation.")

        elif choice == "4":
            # Toggle web search
            rag.use_web_search = not rag.use_web_search
            print(f"\nWeb search is now {'ON' if rag.use_web_search else 'OFF'}")

            if rag.use_web_search and not rag.brave_api_key:
                print("Warning: Brave API key is not set. Web search will not work.")
                brave_key = input("Enter Brave API key (or press Enter to skip): ").strip()
                if brave_key:
                    rag.brave_api_key = brave_key
                    print("API key set!")
                else:
                    print("No API key provided. Web search remains disabled.")
                    rag.use_web_search = False

        elif choice == "5":
            # Memory management submenu
            while True:
                print("\nMemory Settings:")
                print(
                    f"1. {'Disable' if rag.memory_enabled else 'Enable'} conversation memory (currently {'ON' if rag.memory_enabled else 'OFF'})")
                print(f"2. Clear conversation memory ({len(rag.conversation_memory)} conversations stored)")
                print(f"3. Set memory turns (currently: {rag.memory_turns})")
                print("4. View current conversation memory")
                print("5. Return to main menu")

                memory_choice = input("\nEnter choice (1-5): ").strip()

                if memory_choice == "1":
                    rag.memory_enabled = not rag.memory_enabled
                    print(f"Conversation memory is now {'enabled' if rag.memory_enabled else 'disabled'}")

                elif memory_choice == "2":
                    confirm = input("Are you sure you want to clear conversation memory? (y/n): ").lower().strip()
                    if confirm == "y":
                        result = rag.clear_memory()
                        print(result)

                elif memory_choice == "3":
                    turns = input(
                        f"Enter number of conversation turns to remember (current: {rag.memory_turns}): ").strip()
                    try:
                        rag.memory_turns = int(turns)
                        print(f"Will remember last {rag.memory_turns} conversation turns")
                    except ValueError:
                        print("Invalid number. Memory turns not updated.")

                elif memory_choice == "4":
                    if not rag.conversation_memory:
                        print("No conversation memory stored.")
                    else:
                        print("\nCurrent conversation memory:")
                        for i, (q, a) in enumerate(rag.conversation_memory):
                            print(f"\nConversation {i + 1}:")
                            print(f"User: {q}")
                            print(f"Assistant: {a[:100]}...")

                elif memory_choice == "5":
                    break

                else:
                    print("Invalid choice")

        elif choice == "6":
            # Toggle query expansion
            rag.query_expansion_enabled = not rag.query_expansion_enabled
            status = "enabled" if rag.query_expansion_enabled else "disabled"
            print(f"\nQuery expansion is now {status}")

        elif choice == "7":
            print("\nCurrent search parameters:")
            print(f"Search limit: {rag.search_limit}")
            print(f"Dense weight: {rag.dense_weight:.2f}")
            print(f"Sparse weight: {rag.sparse_weight:.2f}")
            print(f"Temperature: {rag.temperature:.2f}")
            print(f"Max tokens: {rag.max_tokens}")
            print(f"Web search: {'Enabled' if rag.use_web_search else 'Disabled'}")
            print(f"Memory: {'Enabled' if rag.memory_enabled else 'Disabled'}")
            print(f"Query expansion: {'Enabled' if rag.query_expansion_enabled else 'Disabled'}")

            try:
                new_limit = input("\nEnter new search limit (or press Enter to keep current): ").strip()
                if new_limit:
                    rag.search_limit = int(new_limit)

                new_dense = input("Enter new dense weight (or press Enter to keep current): ").strip()
                new_sparse = input("Enter new sparse weight (or press Enter to keep current): ").strip()

                if new_dense or new_sparse:
                    dense = float(new_dense) if new_dense else rag.dense_weight
                    sparse = float(new_sparse) if new_sparse else rag.sparse_weight

                    # Normalize weights
                    total = dense + sparse
                    rag.dense_weight = dense / total
                    rag.sparse_weight = sparse / total

                new_temp = input("Enter new temperature (0.0-1.0, or press Enter to keep current): ").strip()
                if new_temp:
                    rag.temperature = float(new_temp)

                new_max = input("Enter new max tokens (or press Enter to keep current): ").strip()
                if new_max:
                    rag.max_tokens = int(new_max)

                print("\nParameters updated!")
                print(f"New search limit: {rag.search_limit}")
                print(f"New dense weight: {rag.dense_weight:.2f}")
                print(f"New sparse weight: {rag.sparse_weight:.2f}")
                print(f"New temperature: {rag.temperature:.2f}")
                print(f"New max tokens: {rag.max_tokens}")

            except ValueError as e:
                print(f"Error updating parameters: {e}")
                print("Parameters not updated.")

        elif choice == "8":
            print("\nExiting RAG system. Goodbye!")
            break

        else:
            print("\nInvalid choice. Please enter a number between 1 and 8.")

if __name__ == "__main__":
    main()