Update rag_engine.py

rag_engine.py

@@ -1,180 +1,180 @@
-# rag_engine.py
-
-import os
-import warnings
-from typing import List
-from dotenv import load_dotenv
-
-# Suppress warnings
-warnings.filterwarnings('ignore')
-
-# Set environment variables before imports
-os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
-os.environ["USER_AGENT"] = "PythonicAI/1.0 (+adinafraz01@gmail.com)"
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-# Load env variables
-load_dotenv()
-openai_key = os.getenv("OPENAI_API_KEY")
-
-if not openai_key:
-    raise ValueError("OPENAI_API_KEY not found in environment variables. Please create a .env file with your API key.")
-
-os.environ["OPENAI_API_KEY"] = openai_key
-
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_openai import OpenAIEmbeddings, ChatOpenAI
-from langchain_community.vectorstores import FAISS
-from langchain_community.document_loaders import WebBaseLoader
-from langchain_core.prompts import PromptTemplate
-from langchain.schema import Document
-
-# ------------------------------------------
-# DOCUMENT LOADING (URL ONLY)
-# ------------------------------------------
-
-def load_user_document(url: str) -> List[Document]:
-    """Load URL into LangChain documents."""
-    try:
-        if not url:
-            raise ValueError("No URL provided")
-        
-        # Clean URL
-        if not url.startswith(('http://', 'https://')):
-            url = f'https://{url}'
-            
-        # Use WebBaseLoader with custom headers
-        headers = {
-            'User-Agent': os.environ.get("USER_AGENT", "PythonicAI/1.0")
-        }
-        
-        loader = WebBaseLoader(
-            web_paths=[url],
-            requests_per_second=1,
-            requests_kwargs={'headers': headers}
-        )
-        docs = loader.load()
-        
-        if not docs:
-            raise ValueError("No content found in the URL")
-            
-        return docs
-        
-    except Exception as e:
-        raise Exception(f"Error loading URL: {str(e)}")
-
-# ------------------------------------------
-# BUILD RETRIEVER FROM DOCUMENTS
-# ------------------------------------------
-
-def build_retriever_from_docs(docs: List[Document]):
-    """Build FAISS vector store from documents."""
-    if not docs:
-        raise ValueError("No documents provided")
-    
-    # Use the same splitter settings as your example
-    splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200
-    )
-    
-    chunks = splitter.split_documents(docs)
-    
-    # Use the same embedding model as your example
-    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
-    
-    # Build vector store
-    vector_store = FAISS.from_documents(chunks, embeddings)
-    
-    return vector_store.as_retriever(
-        search_type="similarity",
-        search_kwargs={"k": 6}
-    )
-
-# ------------------------------------------
-# PYTHON DOMAIN CHECK
-# ------------------------------------------
-
-_llm = None
-
-def get_llm():
-    """Singleton pattern for LLM to avoid re-initialization."""
-    global _llm
-    if _llm is None:
-        _llm = ChatOpenAI(
-            model="gpt-4o-mini",
-            temperature=0.2,
-            max_tokens=1000,
-            request_timeout=60
-        )
-    return _llm
-
-def is_python_related(question: str) -> bool:
-    """Check if question is Python-related."""
-    python_keywords = [
-        'python', 'numpy', 'pandas', 'matplotlib', 'seaborn',
-        'machine learning', 'deep learning', 'tensorflow', 'pytorch',
-        'scikit-learn', 'sklearn', 'dataframe', 'array', 'plot',
-        'import', 'def ', 'class ', 'function', 'method', 'library',
-        'package', 'module', 'pip install', 'conda install', 'jupyter',
-        'data', 'science', 'analysis', 'visualization', 'programming',
-        'code', 'script', 'algorithm', 'model', 'training', 'prediction'
-    ]
-    
-    # Quick keyword check first
-    question_lower = question.lower()
-    for keyword in python_keywords:
-        if keyword in question_lower:
-            return True
-    
-    return False
-
-# ------------------------------------------
-# ANSWER QUESTION
-# ------------------------------------------
-
-# Define the prompt template (same as your example)
-prompt = PromptTemplate(
-    template="""
-    You are a helpful python assistant. 
-    Answer ONLY from the provided text context.
-    Display answer in a beautiful and user understandable format with proper formatting.
-    Use bullet points, numbered lists, or code blocks when appropriate.
-    If the context is insufficient, just say you don't know.
-
-    Context: {context}
-    Question: {question}
-    """,
-    input_variables=['context', 'question']
-)
-
-def answer_question(question: str, retriever):
-    """Answer question using RAG pipeline."""
-    # Validate topic
-    if not is_python_related(question):
-        return "I can only answer Python and data science related questions based on the provided documentation."
-    
-    try:
-        # Retrieve relevant documents
-        docs = retriever.invoke(question)
-        
-        if not docs:
-            return "I couldn't find relevant information in the documentation to answer this question."
-        
-        # Combine context from retrieved documents
-        context = "\n\n".join(doc.page_content for doc in docs)
-        
-        # Create final prompt
-        final_prompt = prompt.invoke({
-            "context": context,
-            "question": question
-        })
-        
-        # Get answer from LLM
-        llm = get_llm()
-        answer = llm.invoke(final_prompt)
-        
-        return answer.content
-        
-    except Exception as e:
+# rag_engine.py
+
+import os
+import warnings
+from typing import List
+from dotenv import load_dotenv
+
+# Suppress warnings
+warnings.filterwarnings('ignore')
+
+# Set environment variables before imports
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+os.environ["USER_AGENT"] = "PythonicAI/1.0 (+adinafraz01@gmail.com)"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# Load env variables
+load_dotenv()
+openai_key = os.getenv("OPENAI_API_KEY")
+
+if not openai_key:
+    raise ValueError("OPENAI_API_KEY not found in environment variables. Please create a .env file with your API key.")
+
+os.environ["OPENAI_API_KEY"] = openai_key
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_community.vectorstores import FAISS
+from langchain_community.document_loaders import WebBaseLoader
+from langchain_core.prompts import PromptTemplate
+from langchain.schema import Document
+
+# ------------------------------------------
+# DOCUMENT LOADING (URL ONLY)
+# ------------------------------------------
+
+def load_user_document(url: str) -> List[Document]:
+    """Load URL into LangChain documents."""
+    try:
+        if not url:
+            raise ValueError("No URL provided")
+        
+        # Clean URL
+        if not url.startswith(('http://', 'https://')):
+            url = f'https://{url}'
+            
+        # Use WebBaseLoader with custom headers
+        headers = {
+            'User-Agent': os.environ.get("USER_AGENT", "PythonicAI/1.0")
+        }
+        
+        loader = WebBaseLoader(
+            web_paths=[url],
+            requests_per_second=1,
+            requests_kwargs={'headers': headers}
+        )
+        docs = loader.load()
+        
+        if not docs:
+            raise ValueError("No content found in the URL")
+            
+        return docs
+        
+    except Exception as e:
+        raise Exception(f"Error loading URL: {str(e)}")
+
+# ------------------------------------------
+# BUILD RETRIEVER FROM DOCUMENTS
+# ------------------------------------------
+
+def build_retriever_from_docs(docs: List[Document]):
+    """Build FAISS vector store from documents."""
+    if not docs:
+        raise ValueError("No documents provided")
+    
+    # Use the same splitter settings as your example
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200
+    )
+    
+    chunks = splitter.split_documents(docs)
+    
+    # Use the same embedding model as your example
+    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
+    
+    # Build vector store
+    vector_store = FAISS.from_documents(chunks, embeddings)
+    
+    return vector_store.as_retriever(
+        search_type="similarity",
+        search_kwargs={"k": 6}
+    )
+
+# ------------------------------------------
+# PYTHON DOMAIN CHECK
+# ------------------------------------------
+
+_llm = None
+
+def get_llm():
+    """Singleton pattern for LLM to avoid re-initialization."""
+    global _llm
+    if _llm is None:
+        _llm = ChatOpenAI(
+            model="gpt-4o-mini",
+            temperature=0.2,
+            max_tokens=1000,
+            request_timeout=60
+        )
+    return _llm
+
+def is_python_related(question: str) -> bool:
+    """Check if question is Python-related."""
+    python_keywords = [
+        'python', 'numpy', 'pandas', 'matplotlib', 'seaborn',
+        'machine learning', 'deep learning', 'tensorflow', 'pytorch',
+        'scikit-learn', 'sklearn', 'dataframe', 'array', 'plot',
+        'import', 'def ', 'class ', 'function', 'method', 'library',
+        'package', 'module', 'pip install', 'conda install', 'jupyter',
+        'data', 'science', 'analysis', 'visualization', 'programming',
+        'code', 'script', 'algorithm', 'model', 'training', 'prediction'
+    ]
+    
+    # Quick keyword check first
+    question_lower = question.lower()
+    for keyword in python_keywords:
+        if keyword in question_lower:
+            return True
+    
+    return False
+
+# ------------------------------------------
+# ANSWER QUESTION
+# ------------------------------------------
+
+# Define the prompt template (same as your example)
+prompt = PromptTemplate(
+    template="""
+    You are a helpful python assistant. 
+    Answer ONLY from the provided text context.
+    Display answer in a beautiful and user understandable format with proper formatting.
+    Use bullet points, numbered lists, or code blocks when appropriate.
+    If the context is insufficient, just say you don't know.
+
+    Context: {context}
+    Question: {question}
+    """,
+    input_variables=['context', 'question']
+)
+
+def answer_question(question: str, retriever):
+    """Answer question using RAG pipeline."""
+    # Validate topic
+    if not is_python_related(question):
+        return "I can only answer Python and data science related questions based on the provided documentation."
+    
+    try:
+        # Retrieve relevant documents
+        docs = retriever.invoke(question)
+        
+        if not docs:
+            return "I couldn't find relevant information in the documentation to answer this question."
+        
+        # Combine context from retrieved documents
+        context = "\n\n".join(doc.page_content for doc in docs)
+        
+        # Create final prompt
+        final_prompt = prompt.invoke({
+            "context": context,
+            "question": question
+        })
+        
+        # Get answer from LLM
+        llm = get_llm()
+        answer = llm.invoke(final_prompt)
+        
+        return answer.content
+        
+    except Exception as e:
         return f"Error generating answer: {str(e)}"