queryrun.py

import os
# Fix OpenMP conflict on Windows (must be set before importing numpy/faiss)
os.environ.setdefault('KMP_DUPLICATE_LIB_OK', 'TRUE')

import cohere
import numpy as np
import faiss
import pickle
import traceback # Import traceback for detailed error printing
from dotenv import load_dotenv
from langchain_community.docstore.document import Document
# Corrected import based on the deprecation warning
from langchain_community.docstore.in_memory import InMemoryDocstore
from openai import OpenAI  # OpenRouter API for enhanced responses
# Try to import agent-related modules dynamically to avoid version compatibility issues
AGENT_IMPORTS_AVAILABLE = False
AgentExecutor = None
create_openai_tools_agent = None
Tool = None
ChatPromptTemplate = None
MessagesPlaceholder = None

def _load_agent_imports():
    """Dynamically load agent imports only when needed"""
    global AGENT_IMPORTS_AVAILABLE, AgentExecutor, create_openai_tools_agent, Tool, ChatPromptTemplate, MessagesPlaceholder
    if AGENT_IMPORTS_AVAILABLE:
        return True
    
    try:
        import importlib
        agents_module = importlib.import_module('langchain.agents')
        AgentExecutor = agents_module.AgentExecutor
        create_openai_tools_agent = agents_module.create_openai_tools_agent
        
        tools_module = importlib.import_module('langchain_core.tools')
        Tool = tools_module.Tool
        
        prompts_module = importlib.import_module('langchain.prompts')
        ChatPromptTemplate = prompts_module.ChatPromptTemplate
        MessagesPlaceholder = prompts_module.MessagesPlaceholder
        
        AGENT_IMPORTS_AVAILABLE = True
        return True
    except (ImportError, AttributeError) as e:
        AGENT_IMPORTS_AVAILABLE = False
        print(f"Warning: LangChain agent imports failed: {e}")
        print("Agentic features will be disabled. This may be due to version incompatibility.")
        print("Try: pip install --upgrade 'langchain==0.3.14' 'langchain-core==0.3.0'")
        return False

# Note: We use duckduckgo_search directly instead of langchain_community wrapper
# to avoid dependency issues with langchain_core.memory

# Try to import langchain_openai, but make it optional
try:
    from langchain_openai import ChatOpenAI
    LANGCHAIN_OPENAI_AVAILABLE = True
except ImportError:
    LANGCHAIN_OPENAI_AVAILABLE = False
    print("Warning: langchain_openai not installed. Agentic features will be disabled.")
    print("Install with: pip install langchain-openai")

# Load environment variables
load_dotenv()
cohere_api_key = os.getenv("COHEREAPIKEY")
openrouter_api_key = os.getenv("OPENROUTER_API_KEY")  # OpenRouter API key for enhanced responses

# OpenRouter client initialization (free tier available)
if openrouter_api_key:
    openrouter_client = OpenAI(
        api_key=openrouter_api_key,
        base_url="https://openrouter.ai/api/v1",
        default_headers={
            "HTTP-Referer": "https://github.com/julienserbanescu/julienserbanescu-rag",  # Optional
            "X-Title": "Julien Serbanescu RAG System",  # Optional
        }
    )
    # LangChain LLM for agentic style (only if langchain_openai is available)
    if LANGCHAIN_OPENAI_AVAILABLE:
        langchain_llm = ChatOpenAI(
            model="z-ai/glm-4.5-air:free",  # Free model on OpenRouter
            openai_api_key=openrouter_api_key,
            openai_api_base="https://openrouter.ai/api/v1",
            default_headers={
                "HTTP-Referer": "https://github.com/julienserbanescu/julienserbanescu-rag",
                "X-Title": "Julien Serbanescu RAG System",
            },
            temperature=0.7,
            streaming=False  # Disable streaming to avoid SSE JSON errors with OpenRouter
        )
    else:
        langchain_llm = None
        print("Note: LangChain agentic features disabled (langchain_openai not installed)")
else:
    openrouter_client = None
    langchain_llm = None
    print("Warning: OPENROUTER_API_KEY not found. OpenRouter features will be disabled.")

# Initialize Cohere client
if not cohere_api_key:
    raise ValueError("COHERE_API_KEY not found in environment variables")
co = cohere.Client(cohere_api_key)

# --- Custom Cohere Embeddings Class (for query embedding) ---
class CohereEmbeddingsForQuery:
    EMBED_DIM = 1024  # embed-english-v3.0 always returns 1024-dim vectors

    def __init__(self, client):
        self.client = client
        self.embed_dim = self.EMBED_DIM
        self._cache = {}

    def embed_query(self, text):
        try:
            if not isinstance(text, str):
                try:
                    text = str(text)
                except UnicodeEncodeError:
                    import unicodedata
                    text = unicodedata.normalize('NFKD', str(text))

            if text in self._cache:
                return self._cache[text]

            response = self.client.embed(
                texts=[text],
                model="embed-english-v3.0",
                input_type="search_query"
            )
            if hasattr(response, 'embeddings') and len(response.embeddings) > 0:
                embedding = np.array(response.embeddings[0]).astype('float32')
            else:
                print("Warning: No query embedding found in the response. Returning zero vector.")
                embedding = np.zeros(self.embed_dim, dtype=np.float32)

            if len(self._cache) > 500:
                self._cache.pop(next(iter(self._cache)))
            self._cache[text] = embedding
            return embedding
        except Exception as e:
            print(f"Query embedding error: {e}")
            return np.zeros(self.embed_dim, dtype=np.float32)

# --- FAISS Query System ---
class FAISSQuerySystem:
    def __init__(self, persist_dir='docs/faiss/'):
        self.persist_dir = persist_dir
        self.index = None
        self.documents = [] # List to hold LangChain Document objects
        self.metadata_list = [] # List to hold metadata dictionaries
        self.embedding_function = CohereEmbeddingsForQuery(co) # Use the query-specific class
        self.load_index()

    # stream_chat_completions function commented out - no longer needed without DeepSeek API
    # def stream_chat_completions(self, input_text):
    #     # Ensure input_text is properly encoded as a string
    #     if not isinstance(input_text, str):
    #         try:
    #             input_text = str(input_text)
    #         except UnicodeEncodeError:
    #             # If there's an encoding error, try to normalize the text
    #             import unicodedata
    #             input_text = unicodedata.normalize('NFKD', str(input_text))
                
    #     # DeepSeek API commented out - just return the input text as is
    #     # response = client.chat.completions.create(
    #         # model="deepseek-chat",
    #         # messages=[
    #             # {"role": "system", "content": "Your job is to make text more appealing by adding emojis, formatting, and other enhancements. Do not include any awkward markup though."},
    #             # {"role": "user", "content": input_text},
    #         # ],
    #         # stream=False
    #     # )
    #     # try:
    #         # resp = response.choices[0].message.content.split("\n---")[1]
    #     # except:
    #         # resp = response.choices[0].message.content
    #     # # Extracting just the core content without the extra sections
    #     # resp = resp.replace('**', '')  # Remove bold formatting
    #     # resp = resp.replace('*', '') 
    #     # return resp
        
    #     # For now, just return the input text without DeepSeek processing
    #     return input_text


    def load_index(self):
        """Load the FAISS index and associated document/metadata files"""
        faiss_index_path = os.path.join(self.persist_dir, "index.faiss")
        pkl_path = os.path.join(self.persist_dir, "index.pkl")
        metadata_path = os.path.join(self.persist_dir, "metadata.pkl")

        print(f"Loading FAISS index from: {faiss_index_path}")
        print(f"Loading docstore info from: {pkl_path}")
        print(f"Loading separate metadata from: {metadata_path}")

        if not os.path.exists(faiss_index_path) or not os.path.exists(pkl_path):
             raise FileNotFoundError(f"Required index files (index.faiss, index.pkl) not found in {self.persist_dir}")

        try:
            # 1. Load FAISS index
            self.index = faiss.read_index(faiss_index_path)
            print(f"FAISS index loaded successfully with {self.index.ntotal} vectors.")

            # 2. Load LangChain docstore pickle file
            with open(pkl_path, 'rb') as f:
                try:
                    docstore, index_to_docstore_id = pickle.load(f)
                except (KeyError, AttributeError) as e:
                    print(f"Error loading pickle file: {str(e)}")
                    print("This might be due to a Pydantic version mismatch.")
                    print("Attempting to recreate the index...")
                    # Delete the incompatible files
                    if os.path.exists(faiss_index_path):
                        os.remove(faiss_index_path)
                    if os.path.exists(pkl_path):
                        os.remove(pkl_path)
                    if os.path.exists(metadata_path):
                        os.remove(metadata_path)
                    # Recreate the index
                    from test import main as recreate_index
                    recreate_index()
                    # Try loading again
                    with open(pkl_path, 'rb') as f:
                        docstore, index_to_docstore_id = pickle.load(f)
                except UnicodeDecodeError:
                    print("Unicode decode error when loading pickle file. Attempting to handle special characters...")
                    # Try to handle the Unicode decode error
                    import codecs
                    with codecs.open(pkl_path, 'rb', encoding='utf-8', errors='replace') as f:
                        docstore, index_to_docstore_id = pickle.load(f)

            # Verify the types after loading
            print(f"Docstore object loaded. Type: {type(docstore)}")
            print(f"Index-to-ID mapping loaded. Type: {type(index_to_docstore_id)}")

            # Now this line should work
            if isinstance(index_to_docstore_id, dict):
                 print(f"Mapping contains {len(index_to_docstore_id)} entries.")
            else:
                 # This case should ideally not happen now, but good to have a check
                 raise TypeError(f"Expected index_to_docstore_id to be a dict, but got {type(index_to_docstore_id)}")

            if not isinstance(docstore, InMemoryDocstore):
                 # Add a check for the docstore type too
                 print(f"Warning: Expected docstore to be InMemoryDocstore, but got {type(docstore)}")


            # 3. Reconstruct the list of documents in FAISS index order
            self.documents = []
            num_vectors = self.index.ntotal 

            # Verify consistency
            if num_vectors != len(index_to_docstore_id):
                 print(f"Warning: FAISS index size ({num_vectors}) does not match mapping size ({len(index_to_docstore_id)}). Reconstruction might be incomplete.")

            print("Reconstructing document list...")
            reconstructed_count = 0
            missing_in_mapping = 0
            missing_in_docstore = 0
            # Ensure docstore has the 'search' method needed.
            if not hasattr(docstore, 'search'):
                raise AttributeError(f"Loaded docstore object (type: {type(docstore)}) does not have a 'search' method.")

            for i in range(num_vectors):
                docstore_id = index_to_docstore_id.get(i)
                if docstore_id:
                    # Use the correct method for InMemoryDocstore to retrieve by ID
                    doc = docstore.search(docstore_id)
                    if doc:
                        self.documents.append(doc)
                        reconstructed_count += 1
                    else:
                        print(f"Warning: Document with ID '{docstore_id}' (for FAISS index {i}) not found in the loaded docstore.")
                        missing_in_docstore += 1
                else:
                    print(f"Warning: No docstore ID found in mapping for FAISS index {i}.")
                    missing_in_mapping += 1

            print(f"Successfully reconstructed {reconstructed_count} documents.")
            if missing_in_mapping > 0: print(f"Could not find mapping for {missing_in_mapping} indices.")
            if missing_in_docstore > 0: print(f"Could not find {missing_in_docstore} documents in docstore despite having mapping.")


            # 4. Load the separate metadata list
            if os.path.exists(metadata_path):
                with open(metadata_path, 'rb') as f:
                    self.metadata_list = pickle.load(f)
                print(f"Loaded separate metadata list with {len(self.metadata_list)} entries.")

                if len(self.metadata_list) != len(self.documents):
                    print(f"Warning: Mismatch between reconstructed documents ({len(self.documents)}) and loaded metadata list ({len(self.metadata_list)}).")
                    print("Falling back to using metadata attached to Document objects if available.")
                    self.metadata_list = [getattr(doc, 'metadata', {}) for doc in self.documents]
                elif not self.documents and self.metadata_list:
                     print("Warning: Loaded metadata but no documents were reconstructed. Discarding metadata.")
                     self.metadata_list = []

            else:
                print("Warning: Separate metadata file (metadata.pkl) not found.")
                print("Attempting to use metadata attached to Document objects.")
                self.metadata_list = [getattr(doc, 'metadata', {}) for doc in self.documents]

            print(f"Final document count: {len(self.documents)}")
            print(f"Final metadata count: {len(self.metadata_list)}")

        except FileNotFoundError as e:
             print(f"Error loading index files: {e}")
             raise
        except Exception as e:
            print(f"An unexpected error occurred during index loading: {e}")
            traceback.print_exc()
            raise

    def is_interview_style_question(self, query):
        """Detect if the query is an interview-style question that would benefit from enhanced AI response"""
        query_lower = query.lower()
        
        # Interview-style question patterns
        interview_patterns = [
            "tell me about", "can you tell me", "describe", "explain",
            "what makes you", "why did you", "how did you", "what inspired",
            "walk me through", "give me an example", "share a story",
            "what was your role", "what challenges", "what was it like",
            "how do you approach", "what's your experience with",
            "what skills", "what technologies", "what projects",
            "what's your background", "what's your journey",
            "what are your strengths", "what are you passionate about",
            "what motivates you", "what's your philosophy",
            "how would you", "what would you do if",
            "describe a time when", "tell me about a project where"
        ]
        
        # Check for interview patterns
        for pattern in interview_patterns:
            if pattern in query_lower:
                return True
        
        # Check for question words that suggest interview context
        question_words = ["why", "how", "what", "when", "where", "which", "who"]
        if any(query_lower.startswith(word) for word in question_words):
            # Additional context clues for interview questions
            interview_context = [
                "experience", "project", "work", "study", "research", "develop",
                "create", "build", "learn", "achieve", "accomplish", "solve",
                "challenge", "problem", "team", "collaborate", "lead", "manage"
            ]
            if any(context in query_lower for context in interview_context):
                return True
        
        return False

    def preprocess_query(self, query):
        """Preprocess query to improve retrieval and context understanding"""
        if not isinstance(query, str):
            try:
                query = str(query)
            except UnicodeEncodeError:
                import unicodedata
                query = unicodedata.normalize('NFKD', str(query))
        
        # Handle personal references - treat "you" as "Julien" for better context
        query = query.replace("what have you done", "what has Julien done")
        query = query.replace("what do you do", "what does Julien do") 
        query = query.replace("your experience", "Julien's experience")
        query = query.replace("your research", "Julien's research")
        query = query.replace("your projects", "Julien's projects")
        query = query.replace("your background", "Julien's background")
        query = query.replace("you have", "Julien has")
        query = query.replace("you worked", "Julien worked")
        query = query.replace("you studied", "Julien studied")
        query = query.replace("you are", "Julien is")
        query = query.replace("you do", "Julien does")
        
        # Add context keywords for better retrieval
        context_keywords = ["Julien Serbanescu", "portfolio", "projects", "research", "experience", "background"]
        query_lower = query.lower()
        
        # If query doesn't contain personal context, add it
        if not any(keyword in query_lower for keyword in ["julien", "you", "your", "his", "he"]):
            query = f"Julien Serbanescu {query}"
        
        return query

    def search(self, query, k=3):
        """Search the index and return relevant documents with metadata and scores"""
        if not self.index or self.index.ntotal == 0:
            print("Warning: FAISS index is not loaded or is empty.")
            return []
        if not self.documents:
            print("Warning: No documents were successfully loaded.")
            return []

        actual_k = min(k, len(self.documents))
        if actual_k == 0:
             return []

        # Preprocess the query for better retrieval
        processed_query = self.preprocess_query(query)
        print(f"Original query: {query}")
        print(f"Processed query: {processed_query}")

        query_embedding = self.embedding_function.embed_query(processed_query)
        if np.all(query_embedding == 0):
             print("Warning: Query embedding failed, search may be ineffective.")

        query_embedding_batch = np.array([query_embedding])
        distances, indices = self.index.search(query_embedding_batch, actual_k)
        results = []
        retrieved_indices = indices[0]

        for i, idx in enumerate(retrieved_indices):
            if idx == -1:
                continue

            if idx < len(self.documents):
                doc = self.documents[idx]
                metadata = self.metadata_list[idx] if idx < len(self.metadata_list) else getattr(doc, 'metadata', {})
                distance = distances[0][i]
                # Since we're using inner product with normalized vectors, distance is already cosine similarity
                similarity_score = float(distance) if distance > 0 else 0.0

                # Ensure content is properly encoded as a string
                content = getattr(doc, 'page_content', str(doc))
                if not isinstance(content, str):
                    try:
                        content = str(content)
                    except UnicodeEncodeError:
                        # If there's an encoding error, try to normalize the text
                        import unicodedata
                        content = unicodedata.normalize('NFKD', str(content))

                results.append({
                    "content": content,
                    "metadata": metadata,
                    "score": float(similarity_score)
                })
            else:
                print(f"Warning: Search returned index {idx} which is out of bounds for loaded documents ({len(self.documents)}).")

        results.sort(key=lambda x: x['score'], reverse=True)
        return results

    def generate_openrouter_response(self, query, context_docs):
        """Generate response using OpenRouter as the primary response generator"""
        if not openrouter_client:
            print("OpenRouter client not available, falling back to Cohere")
            return self.generate_response(query, context_docs)
        
        if not context_docs:
            try:
                response = openrouter_client.chat.completions.create(
                    model="z-ai/glm-4.5-air:free",  # Free model on OpenRouter
                    messages=[
                        {
                            "role": "system", 
                            "content": "You are Julien Serbanescu, a computer science student and AI researcher. Answer questions about your background, projects, and experience in a professional, engaging manner. Be specific and provide concrete examples when possible. If you don't have specific information, acknowledge this and suggest how the user might rephrase their question."
                        },
                        {
                            "role": "user", 
                            "content": f"I could not find relevant documents in my knowledge base to answer your question: '{query}'. Please provide a general response about your background and suggest how the user might rephrase their question."
                        }
                    ],
                    temperature=0.7,
                    max_tokens=1000
                )
                return response.choices[0].message.content
            except Exception as e:
                error_msg = str(e)
                if "402" in error_msg or "credits" in error_msg.lower() or "Insufficient" in error_msg:
                    print(f"OpenRouter requires credits. Falling back to Cohere...")
                else:
                    print(f"Error calling OpenRouter without documents: {e}")
                # Fallback to Cohere
                return self.generate_response(query, [])

        # Format context documents for OpenRouter
        context_text = ""
        for i, doc in enumerate(context_docs[:5]):  # Limit to top 5 docs
            content = doc['content']
            if not isinstance(content, str):
                try:
                    content = str(content)
                except UnicodeEncodeError:
                    import unicodedata
                    content = unicodedata.normalize('NFKD', str(content))
            
            source = doc['metadata'].get('source', 'Unknown')
            context_text += f"\n--- Source {i+1} ({source}) ---\n{content[:2000]}\n"

        try:
            response = openrouter_client.chat.completions.create(
                model="z-ai/glm-4.5-air:free",  # Free model on OpenRouter
                messages=[
                    {
                        "role": "system", 
                        "content": f"""You are Julien Serbanescu, a computer engineering student and AI researcher. Answer the user's question based on the provided context documents about your background, projects, and experience. 

Context about Julien Serbanescu:
{context_text}

Guidelines:
- Answer as if you are Julien speaking in first person
- Be specific and provide concrete examples from the context
- Use a professional but engaging tone
- If the context doesn't contain enough information, acknowledge this and provide what you can
- Structure your response clearly with specific examples
- Show enthusiasm and passion for your work
- For technical questions, provide detailed explanations
- For general questions, give comprehensive but concise answers"""
                    },
                    {
                        "role": "user", 
                        "content": query
                    }
                ],
                temperature=0.7,
                max_tokens=1500
            )
            return response.choices[0].message.content
        except Exception as e:
            error_msg = str(e)
            # Check if it's a credits/402 error
            if "402" in error_msg or "credits" in error_msg.lower() or "Insufficient" in error_msg:
                print(f"OpenRouter requires credits. Falling back to Cohere...")
            else:
                print(f"Error calling OpenRouter: {e}")
            # Fallback to Cohere
            return self.generate_response(query, context_docs)

    def generate_agentic_response(self, query, context_docs):
        """Generate response using LangChain agentic style with tools including web search
        
        Returns:
            tuple: (response_text, web_search_results) where web_search_results is a list of search results
        """
        # Try to load agent imports if not already loaded
        if not _load_agent_imports():
            print("LangChain agent imports not available, falling back to standard OpenRouter")
            response = self.generate_openrouter_response(query, context_docs)
            return response, []
        
        if not langchain_llm:
            print("LangChain LLM not available, falling back to standard OpenRouter")
            response = self.generate_openrouter_response(query, context_docs)
            return response, []
        
        # Store web search results for adding to sources
        web_search_results = []
        
        # Create tools for the agent
        def search_knowledge_base(search_query: str = None) -> str:
            """Search the knowledge base for relevant information. If no query is provided, uses the original question."""
            # Use original query if no search query provided
            if not search_query or search_query.strip() == "":
                search_query = query
            print(f"\n📚 [AGENT] Knowledge base search tool activated! Searching for: '{search_query}'")
            results = self.search(search_query, k=3)
            if not results:
                print("⚠️ [AGENT] No results found in knowledge base.")
                return "No relevant information found in the knowledge base."
            print(f"✅ [AGENT] Knowledge base search completed. Found {len(results)} results.")
            formatted_results = []
            for i, doc in enumerate(results, 1):
                formatted_results.append(f"Source {i}: {doc['content'][:500]}...")
            return "\n".join(formatted_results)
        
        def get_context_summary(_=None) -> str:
            """Get a summary of available context documents. Takes no arguments (ignores any provided)."""
            if not context_docs:
                return "No context documents available."
            summary = f"Found {len(context_docs)} relevant documents:\n"
            for i, doc in enumerate(context_docs[:5], 1):
                source = doc['metadata'].get('source', 'Unknown')
                summary += f"{i}. {source}\n"
            return summary
        
        def search_web(search_query: str) -> str:
            """Search the web for current information. Use this when you need up-to-date information, recent news, or details not in the knowledge base. Always use this for questions about current events, recent developments, or when the knowledge base doesn't have enough information."""
            print(f"\n🔍 [AGENT] Web search tool activated! Searching for: '{search_query}'")
            try:
                # Use direct DuckDuckGo search (more reliable than langchain wrapper)
                from ddgs import DDGS
                with DDGS() as ddgs:
                    # Get search results with URLs
                    ddg_results = list(ddgs.text(search_query, max_results=5))
                    
                    # Debug: Print first result structure to understand format
                    if ddg_results and len(ddg_results) > 0:
                        print(f"🔍 [DEBUG] First result keys: {list(ddg_results[0].keys()) if isinstance(ddg_results[0], dict) else 'Not a dict'}")
                    
                    # Format results
                    formatted_results = []
                    urls = []
                    for r in ddg_results:
                        if not isinstance(r, dict):
                            continue
                        title = r.get('title', '')
                        body = r.get('body', '')
                        # Try different possible field names for URL
                        href = r.get('href', '') or r.get('url', '') or r.get('link', '') or r.get('href', '')
                        if title or body:
                            formatted_results.append(f"{title}: {body}")
                        if href:
                            urls.append(href)
                        # Also check if the result itself is a dict with URL info
                        if not href:
                            # Try to find any URL-like value in all fields
                            for key, value in r.items():
                                if isinstance(value, str) and ('http://' in value or 'https://' in value):
                                    if value not in urls:
                                        urls.append(value)
                                        print(f"🔍 [DEBUG] Found URL in field '{key}': {value}")
                                        break
                    
                    results = "\n".join(formatted_results) if formatted_results else "No results found."
                    
                print(f"✅ [AGENT] Web search completed. Found {len(ddg_results)} results with {len(urls)} URLs.")
                if urls:
                    print(f"🔍 [DEBUG] URLs found: {urls[:3]}")  # Show first 3 URLs for debugging
                else:
                    print(f"⚠️ [DEBUG] No URLs found in search results!")
                
                # Store search results with metadata for sources
                # Always include urls field, even if empty
                web_search_results.append({
                    'query': search_query,
                    'content': results,
                    'type': 'web_search',
                    'source': f'Web Search: {search_query}',
                    'urls': urls[:5] if urls else []  # Store up to 5 URLs, empty list if none
                })
                
                # Debug: Verify URLs are stored
                print(f"🔍 [DEBUG] Stored web_search_result with {len(web_search_results[-1].get('urls', []))} URLs")
                
                # Format results with URLs if found
                if urls:
                    url_list = "\n".join([f"  - {url}" for url in urls[:5]])
                    return f"Web search results for '{search_query}':\n{results}\n\nRelevant URLs:\n{url_list}"
                else:
                    return f"Web search results for '{search_query}':\n{results}"
            except Exception as e:
                print(f"Error in web search: {e}")
                traceback.print_exc()
                return f"Could not perform web search: {str(e)}"
        
        # Define tools (only if Tool class is available)
        if Tool is None:
            print("Tool class not available, falling back to standard OpenRouter")
            response = self.generate_openrouter_response(query, context_docs)
            return response, []
        
        tools = [
            Tool(
                name="search_knowledge_base",
                func=search_knowledge_base,
                description="Search the knowledge base for information about Julien Serbanescu's projects, experience, research, and background. Takes a search query string as input (or uses the original question if no query provided). Use this when you need to find specific information from the indexed documents."
            ),
            Tool(
                name="search_web",
                func=search_web,
                description="Search the internet for current, up-to-date information. Takes a search query string as input (or uses the original question if no query provided). Use this when: 1) The knowledge base doesn't have enough information, 2) You need recent news or developments, 3) You need to verify current facts, or 4) The question is about something that might have changed recently. Only use this tool when it would add value to the response."
            ),
            Tool(
                name="get_context_summary",
                func=get_context_summary,
                description="Get a summary of the currently available context documents. Takes no arguments. Use this to understand what information is available."
            )
        ]
        
        # Format context for the agent
        context_text = ""
        if context_docs:
            for i, doc in enumerate(context_docs[:5]):
                content = doc['content']
                if not isinstance(content, str):
                    try:
                        content = str(content)
                    except UnicodeEncodeError:
                        import unicodedata
                        content = unicodedata.normalize('NFKD', str(content))
                source = doc['metadata'].get('source', 'Unknown')
                context_text += f"\n--- Source {i+1} ({source}) ---\n{content[:2000]}\n"
        
        # Check if prompt classes are available
        if ChatPromptTemplate is None or MessagesPlaceholder is None:
            print("Prompt classes not available, falling back to standard OpenRouter")
            response = self.generate_openrouter_response(query, context_docs)
            return response, []
        
        # Create agent prompt with context included in system message
        # Escape curly braces in context_text to avoid format string errors with ChatPromptTemplate
        context_display = context_text if context_text else "No initial context provided. Use the search_knowledge_base tool to find information."
        # Double all braces to escape them for ChatPromptTemplate
        context_display = context_display.replace("{", "{{").replace("}", "}}")
        
        system_message = f"""You are an AI assistant representing Julien Serbanescu, a computer engineering student and AI researcher. 
You have access to tools that can search his knowledge base, search the web, and retrieve information about his projects, experience, and background.

Available Context:
{context_display}

Guidelines:
- Answer as if you are Julien speaking in first person
- Workflow: 1) First review the available context from the knowledge base, 2) Use search_web tool if you need current information, recent developments, or additional context not in the knowledge base, 3) Synthesize information from both sources in your answer
- Use search_web when it would add value: for current events, recent updates, or when knowledge base information seems incomplete
- Be specific and provide concrete examples from the context
- Use a professional but engaging tone
- Structure your response clearly with specific examples
- Show enthusiasm and passion for the work
- For technical questions, provide detailed explanations
- For general questions, give comprehensive but concise answers
- IMPORTANT: When using web search results, ALWAYS include the URLs in your response so users know you used web search. Format URLs clearly, e.g., "According to [Source Name](URL)..." or list them at the end as "Sources: URL1, URL2, ..."
- Make it clear when information comes from web search vs knowledge base"""
        
        prompt = ChatPromptTemplate.from_messages([
            ("system", system_message),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])
        
        try:
            # Check if agent classes are available
            if not AGENT_IMPORTS_AVAILABLE or create_openai_tools_agent is None or AgentExecutor is None:
                print("Agent classes not available, falling back to standard OpenRouter")
                response = self.generate_openrouter_response(query, context_docs)
                return response, []
            
            # Create agent
            agent = create_openai_tools_agent(langchain_llm, tools, prompt)
            # max_iterations=5 allows for: knowledge base search → web search → final answer (with some flexibility)
            # Each iteration: LLM thinks → calls tool → LLM processes result → repeat
            agent_executor = AgentExecutor(
                agent=agent, 
                tools=tools, 
                verbose=True, 
                max_iterations=3,
                max_execution_time=45,
                return_intermediate_steps=True,
                early_stopping_method="generate"
            )
            
            # Run the agent
            print("\n🤖 [AGENT] Starting agentic response generation...")
            try:
                result = agent_executor.invoke({
                    "input": query
                })
                response_text = result.get("output", "")
                
                # If agent has output, use it (even if it hit max iterations)
                if response_text and response_text.strip():
                    # Remove "Agent stopped" message if present, but keep the actual response
                    if "Agent stopped due to max iterations" in response_text:
                        # Clean up the stop message but keep the response content
                        response_text = response_text.replace("Agent stopped due to max iterations.", "").strip()
                        if response_text:
                            print("\n⚠️ [AGENT] Hit max iterations, but using partial response from agent.")
                        else:
                            # Only had the stop message, no actual content
                            print("\n⚠️ [AGENT] Hit max iterations with no useful output. Falling back...")
                            response = self.generate_openrouter_response(query, context_docs)
                            return response, web_search_results
                    else:
                        # Has response text, use it
                        print("\n✅ [AGENT] Agent completed successfully.")
                else:
                    # No response text at all, fall back
                    print("\n⚠️ [AGENT] No response from agent. Falling back to standard response...")
                    response = self.generate_openrouter_response(query, context_docs)
                    return response, web_search_results
                
            except Exception as agent_error:
                # Handle streaming errors or other agent issues
                error_str = str(agent_error)
                if "SSE" in error_str or "stream" in error_str.lower() or "JSON error" in error_str:
                    print(f"\n⚠️ [AGENT] Streaming error detected: {error_str[:100]}")
                    print("   Falling back to standard OpenRouter response (non-agentic)...")
                    response = self.generate_openrouter_response(query, context_docs)
                    return response, []
                elif "max iterations" in error_str.lower() or "Agent stopped" in error_str:
                    print(f"\n⚠️ [AGENT] Hit iteration limit: {error_str[:100]}")
                    print("   Falling back to standard OpenRouter response (non-agentic)...")
                    response = self.generate_openrouter_response(query, context_docs)
                    return response, web_search_results if web_search_results else []
                else:
                    # Re-raise if it's a different error
                    raise
            
            # Show summary of tool usage
            if web_search_results:
                print(f"\n🌐 [AGENT] Web search was used! Found {len(web_search_results)} web search result(s).")
                for i, web_result in enumerate(web_search_results, 1):
                    print(f"   Web Search {i}: '{web_result.get('query', 'Unknown')}' - {len(web_result.get('urls', []))} URLs found")
                
                # Append web search URLs to response so users know web search was used
                all_urls = []
                for web_result in web_search_results:
                    urls = web_result.get('urls', [])
                    all_urls.extend(urls)
                
                if all_urls:
                    # Remove duplicates while preserving order
                    seen = set()
                    unique_urls = []
                    for url in all_urls:
                        if url not in seen:
                            seen.add(url)
                            unique_urls.append(url)
                    
                    # Append URLs to response
                    if unique_urls:
                        url_section = "\n\n--- Web Search Sources ---\n"
                        url_section += "This response used web search to find current information. Sources:\n"
                        for i, url in enumerate(unique_urls[:5], 1):  # Show up to 5 URLs
                            url_section += f"{i}. {url}\n"
                        response_text += url_section
            else:
                print("\n📋 [AGENT] Only knowledge base was used (no web search needed).")
            
            return response_text, web_search_results
        except Exception as e:
            print(f"Error in agentic response: {e}")
            traceback.print_exc()
            # Fallback to standard OpenRouter
            response = self.generate_openrouter_response(query, context_docs)
            return response, []

    def generate_hybrid_response(self, query, context_docs, use_agent=False):
        """Generate response using OpenRouter (fast) or agentic style (thorough), with Cohere fallback.
        
        Args:
            use_agent: If True, use the multi-step LangChain agent (slower, can do web search).
                       If False (default), use a single OpenRouter LLM call (much faster).
        
        Returns:
            tuple: (response_text, web_search_results) where web_search_results is a list of search results
        """
        if use_agent and _load_agent_imports() and langchain_llm:
            print("Using LangChain agentic style for enhanced response...")
            return self.generate_agentic_response(query, context_docs)
        elif openrouter_client:
            print("Using OpenRouter for enhanced response...")
            response = self.generate_openrouter_response(query, context_docs)
            return response, []
        else:
            print("OpenRouter not available, using Cohere fallback...")
            response = self.generate_response(query, context_docs)
            return response, []

    def generate_response(self, query, context_docs):
        """Generate RAG response using Cohere's chat API"""
        if not context_docs:
            print("No context documents provided to generate_response.")
            try:
                response = co.chat(
                    message=f"I could not find relevant documents in my knowledge base to answer your question: '{query}'. Please try rephrasing or asking about topics covered in the source material.",
                    model="command-r7b-12-2024",  # Updated from command-r (deprecated)
                    temperature=0.3,
                    preamble="You are an AI assistant explaining limitations."
                )
                return response.text
            except Exception as e:
                 print(f"Error calling Cohere even without documents: {e}")
                 return "I could not find relevant documents and encountered an error trying to respond."

        formatted_docs = []
        # Process documents in batches to reduce memory usage
        batch_size = 3
        for i in range(0, len(context_docs), batch_size):
            batch_end = min(i + batch_size, len(context_docs))
            for j in range(i, batch_end):
                doc = context_docs[j]
                # Ensure content is properly encoded as a string
                content = doc['content']
                if not isinstance(content, str):
                    try:
                        content = str(content)
                    except UnicodeEncodeError:
                        # If there's an encoding error, try to normalize the text
                        import unicodedata
                        content = unicodedata.normalize('NFKD', str(content))
                
                content_preview = content[:3000]
                doc_info = f"Source: {doc['metadata'].get('source', 'Unknown')}\n"
                doc_info += f"Type: {doc['metadata'].get('type', 'Unknown')}\n"
                doc_info += f"Content Snippet: {content_preview}"
                formatted_docs.append({"title": f"Document {j+1} (Source: {doc['metadata'].get('source', 'Unknown')})", "snippet": doc_info})
            
            # Force garbage collection after each batch
            import gc
            gc.collect()

        try:
            response = co.chat(
                message=query,
                documents=formatted_docs,
                model="command-r7b-12-2024",  # Updated from command-r (deprecated)
                temperature=0.3,
                prompt_truncation='AUTO',
                preamble="You are an expert AI assistant helping users learn about Julien Serbanescu's background, projects, and experience. Answer the user's question based on the provided document snippets. When the user asks about 'you' or 'your', they are referring to Julien Serbanescu. Use the document information to provide comprehensive, accurate responses. Cite the source document number (e.g., [Document 1]) when using information from it. If the answer isn't in the documents, state that clearly and suggest what information might be available."
            )
            return response.text
        except Exception as e:
             print(f"Error during Cohere chat API call: {e}")
             traceback.print_exc()
             return "Sorry, I encountered an error while trying to generate a response using the retrieved documents."

def main():
    try:
        # Initialize query system
        query_system = FAISSQuerySystem() # Defaults to 'docs/faiss/'

        # Interactive query loop
        print("\n--- FAISS RAG Query System ---")
        print("Ask questions about the content indexed from web, PDFs, and audio.")
        print("Type 'exit' or 'quit' to stop.")

        while True:
            query = input("\nYour question: ")
            if query.lower() in ('exit', 'quit'):
                print("Exiting...")
                break
            if not query:
                continue

            try:
                # 1. Search for relevant documents
                print("Searching for relevant documents...")
                docs = query_system.search(query, k=8) # Get top 8 results for better context

                if not docs:
                    print("Could not find relevant documents in the knowledge base.")
                    response, web_results = query_system.generate_hybrid_response(query, [], use_agent=True)
                    print("\nResponse:")
                    print("-" * 50)
                    print(response)
                    print("-" * 50)
                    
                    # Show web search sources if any
                    if web_results:
                        print("\nWeb Search Sources:")
                        for i, web_result in enumerate(web_results, 1):
                            print(f"\n--- Web Source {i} ---")
                            print(f"  Query: {web_result.get('query', 'Unknown')}")
                            print(f"  Type: {web_result.get('type', 'web_search')}")
                            print(f"  Content: {web_result.get('content', '')[:250]}...")
                    continue

                print(f"Found {len(docs)} relevant document chunks.")

                # 2. Generate and display response using hybrid RAG
                print("Generating response based on documents...")
                response, web_results = query_system.generate_hybrid_response(query, docs, use_agent=True)
                print("\nResponse:")
                print("-" * 50)
                print(response)
                print("-" * 50)

                # 3. Show all sources (knowledge base + web search combined)
                print("\nRetrieved Sources (Snippets):")
                source_num = 1
                
                # Show knowledge base sources
                for i, doc in enumerate(docs, 1):
                    print(f"\n--- Source {source_num} ---")
                    print(f"  Score: {doc['score']:.4f}")
                    print(f"  Source File: {doc['metadata'].get('source', 'Unknown')}")
                    print(f"  Type: {doc['metadata'].get('type', 'Unknown')}")
                    if 'page' in doc['metadata']:
                         print(f"  Page (PDF): {doc['metadata']['page']}")
                    print(f"  Content: {doc['content'][:250]}...")
                    source_num += 1
                
                # Show web search sources integrated into the list
                if web_results:
                    for i, web_result in enumerate(web_results, 1):
                        print(f"\n--- Source {source_num} (Web Search) ---")
                        print(f"  Score: N/A (Web Search)")
                        print(f"  Query: {web_result.get('query', 'Unknown')}")
                        print(f"  Type: {web_result.get('type', 'web_search')}")
                        urls = web_result.get('urls', [])
                        if urls:
                            print(f"  URLs: {', '.join(urls[:3])}")  # Show first 3 URLs
                        content = web_result.get('content', '')
                        if content:
                            print(f"  Content: {content[:250]}...")
                        source_num += 1

            except Exception as e:
                print(f"\nAn error occurred while processing your query: {e}")
                traceback.print_exc()

    except FileNotFoundError as e:
         print(f"\nInitialization Error: Could not find necessary index files.")
         print(f"Details: {e}")
         print("Please ensure you have run the indexing script first and the 'docs/faiss/' directory contains 'index.faiss' and 'index.pkl'.")
    except Exception as e:
        print(f"\nA critical initialization error occurred: {e}")
        traceback.print_exc()

if __name__ == "__main__":
    main()