""" 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)