Spaces:
Sleeping
Sleeping
| 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", {}) | |