Spaces:
Running
Running
| # Utilities to build a RAG system to query information from the | |
| # gwIAS search pipeline using Langchain | |
| # Thanks to Pablo Villanueva Domingo for sharing his CAMELS template | |
| # https://huggingface.co/spaces/PabloVD/CAMELSDocBot | |
| from langchain import hub | |
| from langchain_chroma import Chroma | |
| from langchain_core.output_parsers import StrOutputParser | |
| from langchain_core.runnables import RunnablePassthrough | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from langchain_community.document_loaders import WebBaseLoader | |
| from langchain.schema import Document | |
| import requests | |
| import json | |
| import base64 | |
| from bs4 import BeautifulSoup | |
| import re | |
| def github_to_raw(url): | |
| """Convert GitHub URL to raw content URL""" | |
| return url.replace("github.com", "raw.githubusercontent.com").replace("/blob/", "/") | |
| def load_github_notebook(url): | |
| """Load Jupyter notebook from GitHub URL using GitHub API""" | |
| try: | |
| # Convert GitHub blob URL to API URL | |
| if "github.com" in url and "/blob/" in url: | |
| # Extract owner, repo, branch and path from URL | |
| parts = url.replace("https://github.com/", "").split("/") | |
| owner = parts[0] | |
| repo = parts[1] | |
| branch = parts[3] # usually 'main' or 'master' | |
| path = "/".join(parts[4:]) | |
| api_url = f"https://api.github.com/repos/{owner}/{repo}/contents/{path}?ref={branch}" | |
| else: | |
| raise ValueError("URL must be a GitHub blob URL") | |
| # Fetch notebook content | |
| response = requests.get(api_url) | |
| response.raise_for_status() | |
| content_data = response.json() | |
| if content_data.get('encoding') == 'base64': | |
| notebook_content = base64.b64decode(content_data['content']).decode('utf-8') | |
| else: | |
| notebook_content = content_data['content'] | |
| # Parse notebook JSON | |
| notebook = json.loads(notebook_content) | |
| docs = [] | |
| cell_count = 0 | |
| # Process each cell | |
| for cell in notebook.get('cells', []): | |
| cell_count += 1 | |
| cell_type = cell.get('cell_type', 'unknown') | |
| source = cell.get('source', []) | |
| # Join source lines | |
| if isinstance(source, list): | |
| content = ''.join(source) | |
| else: | |
| content = str(source) | |
| if content.strip(): # Only add non-empty cells | |
| metadata = { | |
| 'source': url, | |
| 'cell_type': cell_type, | |
| 'cell_number': cell_count, | |
| 'name': f"{url} - Cell {cell_count} ({cell_type})" | |
| } | |
| # Add cell type prefix for better context | |
| formatted_content = f"[{cell_type.upper()} CELL {cell_count}]\n{content}" | |
| docs.append(Document(page_content=formatted_content, metadata=metadata)) | |
| return docs | |
| except Exception as e: | |
| print(f"Error loading notebook from {url}: {str(e)}") | |
| return [] | |
| def clean_text(text): | |
| """Clean text content from a webpage""" | |
| # Remove excessive newlines | |
| text = re.sub(r'\n{3,}', '\n\n', text) | |
| # Remove excessive whitespace | |
| text = re.sub(r'\s{2,}', ' ', text) | |
| return text.strip() | |
| def clean_github_content(html_content): | |
| """Extract meaningful content from GitHub pages""" | |
| # Ensure we're working with a BeautifulSoup object | |
| if isinstance(html_content, str): | |
| soup = BeautifulSoup(html_content, 'html.parser') | |
| else: | |
| soup = html_content | |
| # Remove navigation, footer, and other boilerplate | |
| for element in soup.find_all(['nav', 'footer', 'header']): | |
| element.decompose() | |
| # For README and code files | |
| readme_content = soup.find('article', class_='markdown-body') | |
| if readme_content: | |
| return clean_text(readme_content.get_text()) | |
| # For code files | |
| code_content = soup.find('table', class_='highlight') | |
| if code_content: | |
| return clean_text(code_content.get_text()) | |
| # For directory listings | |
| file_list = soup.find('div', role='grid') | |
| if file_list: | |
| return clean_text(file_list.get_text()) | |
| # Fallback to main content | |
| main_content = soup.find('main') | |
| if main_content: | |
| return clean_text(main_content.get_text()) | |
| # If no specific content found, get text from body | |
| body = soup.find('body') | |
| if body: | |
| return clean_text(body.get_text()) | |
| # Final fallback | |
| return clean_text(soup.get_text()) | |
| class GitHubLoader(WebBaseLoader): | |
| """Custom loader for GitHub pages with better content cleaning""" | |
| def clean_text(self, text): | |
| """Clean text content""" | |
| # Remove excessive newlines and spaces | |
| text = re.sub(r'\n{2,}', '\n', text) | |
| text = re.sub(r'\s{2,}', ' ', text) | |
| # Remove common GitHub boilerplate | |
| text = re.sub(r'Skip to content|Sign in|Search or jump to|Footer navigation|Terms|Privacy|Security|Status|Docs', '', text) | |
| return text.strip() | |
| def _scrape(self, url: str, *args, **kwargs) -> str: | |
| """Scrape data from URL and clean it. | |
| Args: | |
| url: The URL to scrape | |
| *args: Additional positional arguments | |
| **kwargs: Additional keyword arguments including bs_kwargs | |
| Returns: | |
| str: The cleaned content | |
| """ | |
| response = requests.get(url) | |
| response.raise_for_status() | |
| # For directory listings (tree URLs), use the API | |
| if '/tree/' in url: | |
| # Parse URL components | |
| parts = url.replace("https://github.com/", "").split("/") | |
| owner = parts[0] | |
| repo = parts[1] | |
| branch = parts[3] # usually 'main' or 'master' | |
| path = "/".join(parts[4:]) if len(parts) > 4 else "" | |
| # Construct API URL | |
| api_url = f"https://api.github.com/repos/{owner}/{repo}/contents/{path}?ref={branch}" | |
| api_response = requests.get(api_url) | |
| api_response.raise_for_status() | |
| # Parse directory listing | |
| contents = api_response.json() | |
| if isinstance(contents, list): | |
| # Format directory contents | |
| files = [f"{item['name']} ({item['type']})" for item in contents] | |
| return "Directory contents:\n" + "\n".join(files) | |
| else: | |
| return f"Error: Unexpected API response for {url}" | |
| # For regular files, parse HTML | |
| soup = BeautifulSoup(response.text, 'html.parser') | |
| # For README and markdown files | |
| readme_content = soup.find('article', class_='markdown-body') | |
| if readme_content: | |
| return self.clean_text(readme_content.get_text()) | |
| # For code files | |
| code_content = soup.find('table', class_='highlight') | |
| if code_content: | |
| return self.clean_text(code_content.get_text()) | |
| # For other content, get main content | |
| main_content = soup.find('main') | |
| if main_content: | |
| return self.clean_text(main_content.get_text()) | |
| # Final fallback | |
| return self.clean_text(soup.get_text()) | |
| # Load documentation from urls | |
| def load_docs(): | |
| # Get urls | |
| urlsfile = open("urls.txt") | |
| urls = urlsfile.readlines() | |
| urls = [url.replace("\n","") for url in urls] | |
| urlsfile.close() | |
| # Load documents from URLs | |
| docs = [] | |
| for url in urls: | |
| url = url.strip() | |
| if not url: | |
| continue | |
| # Check if URL is a Jupyter notebook | |
| if url.endswith('.ipynb') and 'github.com' in url and '/blob/' in url: | |
| print(f"Loading notebook: {url}") | |
| notebook_docs = load_github_notebook(url) | |
| docs.extend(notebook_docs) | |
| # Handle Python and Markdown files using raw content | |
| elif url.endswith(('.py', '.md')) and 'github.com' in url and '/blob/' in url: | |
| print(f"Loading raw content: {url}") | |
| try: | |
| raw_url = github_to_raw(url) | |
| loader = WebBaseLoader([raw_url]) | |
| web_docs = loader.load() | |
| # Preserve original URL in metadata | |
| for doc in web_docs: | |
| doc.metadata['source'] = url | |
| docs.extend(web_docs) | |
| except Exception as e: | |
| print(f"Error loading {url}: {str(e)}") | |
| # Handle directory listings | |
| elif '/tree/' in url and 'github.com' in url: | |
| print(f"Loading directory: {url}") | |
| try: | |
| # Parse URL components | |
| parts = url.replace("https://github.com/", "").split("/") | |
| owner = parts[0] | |
| repo = parts[1] | |
| branch = parts[3] # usually 'main' or 'master' | |
| path = "/".join(parts[4:]) if len(parts) > 4 else "" | |
| # Construct API URL | |
| api_url = f"https://api.github.com/repos/{owner}/{repo}/contents/{path}?ref={branch}" | |
| response = requests.get(api_url) | |
| response.raise_for_status() | |
| # Parse directory listing | |
| contents = response.json() | |
| if isinstance(contents, list): | |
| # Format directory contents | |
| content = "Directory contents:\n" + "\n".join([f"{item['name']} ({item['type']})" for item in contents]) | |
| docs.append(Document(page_content=content, metadata={'source': url})) | |
| else: | |
| print(f"Error: Unexpected API response for {url}") | |
| except Exception as e: | |
| print(f"Error loading directory {url}: {str(e)}") | |
| else: | |
| print(f"Loading web page: {url}") | |
| try: | |
| loader = GitHubLoader([url]) # Use custom loader | |
| web_docs = loader.load() | |
| docs.extend(web_docs) | |
| except Exception as e: | |
| print(f"Error loading {url}: {str(e)}") | |
| # Add source URLs as document names for reference | |
| for i, doc in enumerate(docs): | |
| if 'source' in doc.metadata: | |
| doc.metadata['name'] = doc.metadata['source'] | |
| else: | |
| doc.metadata['name'] = f"Document {i+1}" | |
| print(f"Loaded {len(docs)} documents:") | |
| for doc in docs: | |
| print(f" - {doc.metadata.get('name')}") | |
| return docs | |
| def extract_reference(url): | |
| """Extract a reference keyword from the GitHub URL""" | |
| if "blob/main" in url: | |
| return url.split("blob/main/")[-1] | |
| elif "tree/main" in url: | |
| return url.split("tree/main/")[-1] or "root" | |
| return url | |
| # Join content pages for processing | |
| def format_docs(docs): | |
| formatted_docs = [] | |
| for doc in docs: | |
| source = doc.metadata.get('source', 'Unknown source') | |
| reference = f"[{extract_reference(source)}]" | |
| content = doc.page_content | |
| formatted_docs.append(f"{content}\n\nReference: {reference}") | |
| return "\n\n---\n\n".join(formatted_docs) | |
| # Create a RAG chain | |
| def RAG(llm, docs, embeddings): | |
| # Split text | |
| text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) | |
| splits = text_splitter.split_documents(docs) | |
| # Create vector store | |
| vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings) | |
| # Retrieve and generate using the relevant snippets of the documents | |
| retriever = vectorstore.as_retriever() | |
| # Prompt basis example for RAG systems | |
| prompt = hub.pull("rlm/rag-prompt") | |
| # Adding custom instructions to the prompt | |
| template = prompt.messages[0].prompt.template | |
| template_parts = template.split("\nQuestion: {question}") | |
| combined_template = "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, just say that you don't know. "\ | |
| + "Try to keep the answer concise if possible. "\ | |
| + "Write the names of the relevant functions from the retrived code and include code snippets to aid the user's understanding. "\ | |
| + "Include the references used in square brackets at the end of your answer."\ | |
| + template_parts[1] | |
| prompt.messages[0].prompt.template = combined_template | |
| # Create the chain | |
| rag_chain = ( | |
| {"context": retriever | format_docs, "question": RunnablePassthrough()} | |
| | prompt | |
| | llm | |
| | StrOutputParser() | |
| ) | |
| return rag_chain |