import io import zipfile import requests import frontmatter import re import os from dotenv import load_dotenv from openai import OpenAI from minsearch import VectorSearch from sentence_transformers import SentenceTransformer from minsearch import Index from tqdm.auto import tqdm import numpy as np load_dotenv() # Set cache directories for HuggingFace Spaces compatibility os.environ['TRANSFORMERS_CACHE'] = os.environ.get('TRANSFORMERS_CACHE', '/tmp/.cache/huggingface') os.environ['HF_HOME'] = os.environ.get('HF_HOME', '/tmp/.cache/huggingface') os.environ['SENTENCE_TRANSFORMERS_HOME'] = os.environ.get('SENTENCE_TRANSFORMERS_HOME', '/tmp/.cache/sentence-transformers') class RAGChatbot: def __init__(self): self.embedding_model = None self.chunks = None self.embeddings = None self.repo_vindex = None api_key = os.environ.get('OPENROUTER_API_KEY') if not api_key: raise ValueError("OPENROUTER_API_KEY environment variable not set") self.client = OpenAI( base_url="https://openrouter.ai/api/v1", api_key=api_key, ) self.model = "z-ai/glm-4.5-air:free" def initialize(self, repo_owner='Rachel0619', repo_name='PickLLM'): """Initialize the chatbot by loading and processing repository data""" print(f"Initializing RAG chatbot with {repo_owner}/{repo_name}...") data = self._read_repo_data(repo_owner, repo_name) self.chunks = self._process_documents(data) self._build_index() print(f"✅ RAG chatbot initialized with {len(self.chunks)} chunks") def _read_repo_data(self, repo_owner, repo_name): """Read repository data from GitHub""" prefix = 'https://codeload.github.com' url = f'{prefix}/{repo_owner}/{repo_name}/zip/refs/heads/main' resp = requests.get(url) if resp.status_code != 200: raise Exception(f"Failed to download repository: {resp.status_code}") repository_data = [] zf = zipfile.ZipFile(io.BytesIO(resp.content)) for file_info in zf.infolist(): filename = file_info.filename filename_lower = filename.lower() if not (filename_lower.endswith('.md') or filename_lower.endswith('.mdx')): continue try: with zf.open(file_info) as f_in: content = f_in.read().decode('utf-8', errors='ignore') post = frontmatter.loads(content) data = post.to_dict() data['filename'] = filename repository_data.append(data) except Exception as e: print(f'Error processing {filename}: {e}') continue zf.close() return repository_data def _split_markdown_by_level(self, text, level=2): """ Split markdown text by a specific header level. :param text: Markdown text as a string :param level: Header level to split on :return: List of sections as strings """ # This regex matches markdown headers # For level 2, it matches lines starting with "## " header_pattern = r'^(#{' + str(level) + r'} )(.+)$' pattern = re.compile(header_pattern, re.MULTILINE) # Split and keep the headers parts = pattern.split(text) sections = [] for i in range(1, len(parts), 3): # We step by 3 because regex.split() with # capturing groups returns: # [before_match, group1, group2, after_match, ...] # here group1 is "## ", group2 is the header text header = parts[i] + parts[i+1] # "## " + "Title" header = header.strip() # Get the content after this header content = "" if i+2 < len(parts): content = parts[i+2].strip() if content: section = f'{header}\n\n{content}' else: section = header sections.append(section) return sections def _process_documents(self, data): """Process documents into chunks""" chunks = [] for doc in data: doc_copy = doc.copy() doc_content = doc_copy.pop('content', '') if not doc_content: continue sections = self._split_markdown_by_level(doc_content, level=3) for section in sections: section_doc = doc_copy.copy() section_doc['section'] = section chunks.append(section_doc) return chunks def _build_index(self): """Build vector search index from chunks""" print("Building vector index...") self.embedding_model = SentenceTransformer('multi-qa-distilbert-cos-v1') embeddings = [] for d in tqdm(self.chunks): v = self.embedding_model.encode(d['section']) embeddings.append(v) self.embeddings = np.array(embeddings) self.repo_vindex = VectorSearch() self.repo_vindex.fit(self.embeddings, self.chunks) print("✅ Vector index built successfully") def _vector_search(self, query, num_results=2): """Search for relevant context using vector similarity""" if self.repo_vindex is None or self.embedding_model is None: raise Exception("Chatbot not initialized. Call initialize() first.") q = self.embedding_model.encode(query) return self.repo_vindex.search(q, num_results=num_results) def _format_prompt(self, query, context): """Format the prompt for the LLM""" context_text = "\n\n".join([doc.get('section', '') for doc in context]) return f""" You are a friendly assistant for PickLLM, an AI-powered platform that helps users discover and compare Large Language Models. Answer user questions in a conversational, concise way (2-3 sentences max). Be helpful and warm, but keep responses brief and to the point. Question: {query} Relevant context from PickLLM documentation: {context_text} Instructions: - Answer directly and concisely - Use the context provided to give accurate information - If the context doesn't contain the answer, politely say you don't have that specific information - Keep your response under 3 sentences when possible """.strip() def chat(self, query): """ Main chat interface - takes a user query and returns a response """ try: # Search for relevant context context = self._vector_search(query, num_results=2) # Format prompt and call LLM prompt = self._format_prompt(query, context) completion = self.client.chat.completions.create( extra_headers={ "HTTP-Referer": "https://pickllm.com", "X-Title": "PickLLM", }, model=self.model, messages=[ { "role": "system", "content": "You are an AI assistant helping users understand how PickLLM works" }, { "role": "user", "content": prompt } ], temperature=0.7, max_tokens=1000 ) answer = completion.choices[0].message.content return answer except Exception as e: import traceback error_trace = traceback.format_exc() print(f"❌ Error in chat: {e}") print(f"❌ Full traceback:\n{error_trace}") return "I'm sorry, I encountered an error processing your question. Please try again." if __name__ == "__main__": # Test the chatbot chatbot = RAGChatbot() chatbot.initialize('Rachel0619', 'PickLLM') # Test query query = "What dataset are you using to support the recommendation?" print(chatbot._vector_search(query, num_results=1)) response = chatbot.chat(query) print(f"Q: {query}") print(f"A: {response}")