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