import os os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" import pickle import time from pathlib import Path from typing import AsyncIterator, Dict, List, Any, Tuple import httpx from bs4 import BeautifulSoup from dotenv import load_dotenv # LangChain Imports from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_core.vectorstores import InMemoryVectorStore from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter # Load environment variables load_dotenv() # Data persistence directory DATA_DIR = Path("data") DATA_DIR.mkdir(exist_ok=True) DB_PATH = DATA_DIR / "vectorstore.pkl" class RAGManager: def __init__(self): # Retrieve configuration self.api_key = os.getenv("OPENROUTER_API_KEY", "") self.model_name = os.getenv("OPENROUTER_MODEL", "openrouter/owl-alpha") self.embedding_model = os.getenv("EMBEDDING_MODEL", "openai/text-embedding-3-small") self.base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1") if not self.api_key: raise ValueError("OPENROUTER_API_KEY is not set in .env") # Initialize Embeddings pointing to OpenRouter API self.embeddings = OpenAIEmbeddings( model=self.embedding_model, api_key=self.api_key, base_url=self.base_url ) # Initialize Vector Store (In-Memory with custom pickle persistence) self.vector_store = InMemoryVectorStore(self.embeddings) self._load_vector_store() def _load_vector_store(self): """Loads serialized vector store items from disk if available.""" if DB_PATH.exists(): try: with open(DB_PATH, "rb") as f: self.vector_store.store = pickle.load(f) print(f"Loaded existing vector store with {len(self.vector_store.store)} vectors.") except Exception as e: print(f"Error loading vector store pickle: {e}. Starting fresh.") self.vector_store.store = {} else: self.vector_store.store = {} def _save_vector_store(self): """Serializes current vector store items to disk.""" try: with open(DB_PATH, "wb") as f: pickle.dump(self.vector_store.store, f) print(f"Saved vector store with {len(self.vector_store.store)} vectors to disk.") except Exception as e: print(f"Error saving vector store pickle: {e}") def clear_database(self): """Resets the vector database and deletes local pickle files.""" self.vector_store.store = {} if DB_PATH.exists(): try: DB_PATH.unlink() except Exception as e: print(f"Error removing vector store file: {e}") print("Database cleared successfully.") def get_indexed_sources(self) -> List[Dict[str, Any]]: """Summarizes indexed documents in the database.""" sources = {} for doc_id, item in self.vector_store.store.items(): doc = item.get("document") if not doc: continue source_name = doc.metadata.get("source", "Unknown") char_count = len(doc.page_content) timestamp = doc.metadata.get("timestamp", time.time()) if source_name not in sources: sources[source_name] = { "source": source_name, "chunk_count": 0, "total_chars": 0, "added_at": timestamp } sources[source_name]["chunk_count"] += 1 sources[source_name]["total_chars"] += char_count # Keep the oldest timestamp sources[source_name]["added_at"] = min(sources[source_name]["added_at"], timestamp) return list(sources.values()) def add_text_content(self, text: str, source_name: str, chunk_size: int = 800, chunk_overlap: int = 100) -> Dict[str, Any]: """Splits, embeds, and indexes text content.""" if not text.strip(): raise ValueError("Provided text is empty.") # Set up text splitter splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len ) chunks = splitter.split_text(text) if not chunks: raise ValueError("No text chunks generated.") # Wrap in LangChain Document objects now = time.time() docs = [ Document( page_content=chunk, metadata={ "source": source_name, "timestamp": now, "chunk_index": i, "total_chunks": len(chunks) } ) for i, chunk in enumerate(chunks) ] # Add to vector store (handles embedding generation automatically) self.vector_store.add_documents(docs) self._save_vector_store() return { "source": source_name, "chunks_added": len(docs), "total_chars": len(text) } def add_file(self, file_content: bytes, filename: str, chunk_size: int = 800, chunk_overlap: int = 100) -> Dict[str, Any]: """Parses and indexes file bytes depending on file extension.""" ext = Path(filename).suffix.lower() text = "" if ext in (".txt", ".md", ".json", ".xml", ".html"): try: text = file_content.decode("utf-8") except UnicodeDecodeError: # Fallback to latin-1 text = file_content.decode("latin-1") elif ext == ".pdf": from pypdf import PdfReader import io pdf_file = io.BytesIO(file_content) reader = PdfReader(pdf_file) text_parts = [] for i, page in enumerate(reader.pages): page_text = page.extract_text() if page_text: text_parts.append(page_text) text = "\n\n".join(text_parts) if not text.strip(): raise ValueError("Could not extract any text from this PDF file.") else: raise ValueError(f"Unsupported file format: {ext}. Only TXT, MD, and PDF are supported.") return self.add_text_content(text, filename, chunk_size, chunk_overlap) async def add_url(self, url: str, chunk_size: int = 800, chunk_overlap: int = 100) -> Dict[str, Any]: """Scrapes text from a web page and indexes it.""" if not (url.startswith("http://") or url.startswith("https://")): raise ValueError("URL must start with http:// or https://") headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 RAGScraper/1.0" } async with httpx.AsyncClient(follow_redirects=True, timeout=20.0) as client: resp = await client.get(url, headers=headers) resp.raise_for_status() soup = BeautifulSoup(resp.content, "lxml") # Extract page title title = soup.title.string.strip() if soup.title else url # Decompose scripts, styles, and iframe contents to clean text for tag in soup(["script", "style", "noscript", "iframe", "svg", "header", "footer", "nav"]): tag.decompose() # Extract clean visible text content import re raw_text = soup.get_text(separator="\n") # Clean up whitespace cleaned_text = re.sub(r"\n{3,}", "\n\n", raw_text) cleaned_text = re.sub(r"[ \t]+", " ", cleaned_text).strip() if not cleaned_text: raise ValueError("No readable text content found on the webpage.") source_label = f"Web: {title} ({url})" return self.add_text_content(cleaned_text, source_label, chunk_size, chunk_overlap) async def stream_chat( self, query: str, system_prompt: str = None, temperature: float = 0.7, k: int = 4 ) -> AsyncIterator[Tuple[str, Any]]: """ Retrieves matching chunks, builds context, and streams OpenRouter completions. Yields tuples of (event_type, event_data). """ # 1. Check if database is empty if not self.vector_store.store: # Empty database fallback: Chat directly with model without context, but warn the user. context_text = "(No documents loaded. Answering from general knowledge.)" retrieved_data = [] else: # 2. Retrieve top matching documents # similarity_search_with_score returns List[Tuple[Document, float]] # InMemoryVectorStore computes cosine distance where smaller is more similar. # Let's map it to similarity scores (1.0 - distance). matches = self.vector_store.similarity_search_with_score(query, k=k) retrieved_data = [] context_parts = [] for doc, distance in matches: # Cosine distance to similarity: similarity = max(0.0, min(1.0, 1.0 - float(distance))) retrieved_data.append({ "content": doc.page_content, "source": doc.metadata.get("source", "Unknown"), "chunk_index": doc.metadata.get("chunk_index", 0), "total_chunks": doc.metadata.get("total_chunks", 1), "similarity": round(similarity, 4) }) context_parts.append( f"--- SOURCE: {doc.metadata.get('source')} ---\n{doc.page_content}" ) context_text = "\n\n".join(context_parts) # 3. Yield the retrieved chunks first so the frontend can render them immediately yield ("citations", retrieved_data) # 4. Formulate prompts if not system_prompt or not system_prompt.strip(): system_prompt = ( "You are an expert RAG Assistant. Use the retrieved context documents below to provide accurate, " "comprehensive, and helpful answers to the user. Always ground your claims using the sources provided. " "If the context is insufficient or doesn't mention the answer, answer as best as you can based on " "general knowledge, but clearly state that the information was not found in the uploaded documents." ) full_system_content = f"{system_prompt}\n\nRetrieved Context Documents:\n{context_text}" # 5. Initialize ChatOpenAI for Streaming llm = ChatOpenAI( model=self.model_name, api_key=self.api_key, base_url=self.base_url, temperature=temperature, streaming=True ) messages = [ {"role": "system", "content": full_system_content}, {"role": "user", "content": query} ] # 6. Stream tokens try: async for chunk in llm.astream(messages): # chunk is a BaseMessageChunk, its content is the text delta if chunk.content: yield ("token", chunk.content) except Exception as e: yield ("error", f"OpenRouter streaming error: {str(e)}") return yield ("done", {})