src/app_enhanced.py

"""
Enhanced RAG application with crawler integration and modular architecture.
"""

import json
from pathlib import Path
from typing import Optional, List

from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader, WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_ollama import ChatOllama
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.documents import Document

from src.config import CRAWLER_CACHE_PATH, DEFAULT_DOC_DIR

system_prompt = (
    "You are an assistant for question-answering tasks. "
    "Use the following pieces of retrieved context to answer the question. "
    "If you don't know the answer, say that you don't know. "
    "Always cite the source URL if available.\n\n"
    "{context}"
)

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}"),
])


class EnhancedRAGService:
    """Stateful RAG service for local Ollama-based retrieval."""

    def __init__(self, doc_dir: Optional[str] = None, urls: Optional[List[str]] = None):
        self.doc_dir = Path(doc_dir or DEFAULT_DOC_DIR)
        self.urls = urls or []
        self._rag_chain = None
        self._vectorstore = None
        self._llm = ChatOllama(model="llama3", temperature=0)

    @staticmethod
    def format_docs(docs):
        formatted = []
        for doc in docs:
            source = doc.metadata.get("source", "Unknown source")
            formatted.append(f"Source: {source}\n{doc.page_content}")
        return "\n\n---\n\n".join(formatted)

    def _load_pdf_documents(self, doc_dir: Path):
        if not doc_dir.exists():
            return []

        try:
            loader = DirectoryLoader(str(doc_dir), glob="./*.pdf", loader_cls=PyPDFLoader)
            pdf_docs = loader.load()
            if pdf_docs:
                print(f"✓ Loaded {len(pdf_docs)} PDF(s) from {doc_dir}")
            return pdf_docs
        except Exception as exc:
            print(f"✗ Error loading PDFs from {doc_dir}: {exc}")
            return []

    def _load_url_documents(self, urls: Optional[List[str]]):
        docs = []
        if not urls:
            return docs

        print(f"Loading content from {len(urls)} URL(s)...")
        for url in urls:
            try:
                loader = WebBaseLoader(url)
                url_docs = loader.load()
                docs.extend(url_docs)
                print(f"  ✓ Loaded: {url}")
            except Exception as exc:
                print(f"  ✗ Failed to load {url}: {exc}")
        return docs

    def _load_crawler_documents(self):
        crawler_path = Path(CRAWLER_CACHE_PATH)
        if not crawler_path.exists():
            return []

        try:
            with open(crawler_path, "r", encoding="utf-8") as f:
                cached_docs = json.load(f)
            docs = [
                Document(page_content=item["content"], metadata={"source": item.get("url", "Unknown source")})
                for item in cached_docs
                if item.get("content")
            ]
            if docs:
                print(f"✓ Loaded {len(docs)} documents from crawler cache")
            return docs
        except Exception as exc:
            print(f"✗ Error loading crawler cache: {exc}")
            return []

    def load_documents(self, doc_dir: Optional[str] = None, urls: Optional[List[str]] = None):
        doc_path = Path(doc_dir or self.doc_dir)
        docs = []
        docs.extend(self._load_pdf_documents(doc_path))
        docs.extend(self._load_url_documents(urls))
        docs.extend(self._load_crawler_documents())
        return docs

    def build_rag_chain(self, doc_dir: Optional[str] = None, urls: Optional[List[str]] = None):
        docs = self.load_documents(doc_dir, urls)
        if not docs:
            raise ValueError(
                f"No documents found. Add PDFs to '{doc_dir or self.doc_dir}', provide URLs, or run crawler first."
            )

        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = text_splitter.split_documents(docs)
        if not splits:
            raise ValueError("No text content could be extracted from documents.")

        print(f"Processing {len(splits)} text chunks...")
        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
        self._vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)

        retriever = self._vectorstore.as_retriever(search_kwargs={"k": 3})
        self._rag_chain = ({"context": retriever | self.format_docs, "input": RunnablePassthrough()} | prompt | self._llm)
        return self._rag_chain

    def get_rag_chain(self, doc_dir: Optional[str] = None, urls: Optional[List[str]] = None):
        normalized_urls = urls if urls is not None else self.urls
        normalized_urls = normalized_urls or []
        if self._rag_chain is None or normalized_urls != self.urls:
            self.urls = normalized_urls
            self._rag_chain = self.build_rag_chain(doc_dir, normalized_urls)
        return self._rag_chain

    def answer_question(self, question: str, doc_dir: Optional[str] = None, urls: Optional[List[str]] = None) -> str:
        try:
            rag_chain = self.get_rag_chain(doc_dir, urls)
            response = rag_chain.invoke(question)
            return response.content if hasattr(response, "content") else str(response)
        except Exception as exc:
            error_msg = str(exc)
            if "ConnectError" in error_msg or "connection could be made" in error_msg.lower():
                return (
                    "❌ Error: Ollama service is not running.\n\n"
                    "Please start Ollama in a terminal:\n  ollama serve\n\n"
                    "Then ensure the llama3 model is available:\n  ollama pull llama3"
                )
            if "No documents found" in error_msg or "No text content" in error_msg:
                return (
                    "❌ Error: No documents available.\n\n"
                    "Please:\n1. Add PDF files to the 'my_docs/' folder, OR\n"
                    "2. Provide URLs, OR\n"
                    "3. Run the crawler first to crawl a website"
                )
            return f"❌ Error: {error_msg}\n\nPlease check your connection and try again."

    def add_urls(self, url_list: List[str]):
        self.urls = url_list or []
        self._rag_chain = None

    def index_crawler_results(self, documents: List[dict]):
        if not documents:
            raise ValueError("No documents to index")
        with open(CRAWLER_CACHE_PATH, "w", encoding="utf-8") as f:
            json.dump(documents, f, indent=2, ensure_ascii=False)
        self._rag_chain = None
        self._vectorstore = None
        print(f"✓ Indexed {len(documents)} documents from crawler")


service = EnhancedRAGService()


def answer_question(question: str, doc_dir: str = "./my_docs", urls: Optional[List[str]] = None) -> str:
    return service.answer_question(question, doc_dir=doc_dir, urls=urls)


def add_urls(url_list: List[str]):
    return service.add_urls(url_list)


def index_crawler_results(documents: List[dict]):
    return service.index_crawler_results(documents)


def load_documents_from_crawler(doc_dir: Optional[str] = None, urls: Optional[List[str]] = None):
    return service.load_documents(doc_dir=doc_dir, urls=urls)


if __name__ == "__main__":
    answer = answer_question("What are the main features discussed?")
    print("\n--- Answer ---")
    print(answer)