""" RAG Prototype for WebCraft AI ----------------------------- This script demonstrates the "Librarian" (Retrieval) and "Context Manager" architecture. It allows the AI to access ALL 80+ blocks by dynamically retrieving only the relevant ones. """ import os import json import difflib from typing import List, Dict, Any from dotenv import load_dotenv from huggingface_hub import InferenceClient # Import existing schema try: from block_schema import BLOCKLY_SCHEMA except ImportError: # Fallback for standalone run if needed BLOCKLY_SCHEMA = {'blocks': {}} load_dotenv() class RAGLibrarian: """The 'Librarian' that finds relevant blocks.""" def __init__(self, schema: Dict): self.blocks = schema.get('blocks', {}) self.index = self._build_index() def _build_index(self) -> List[Dict]: """Create a searchable text index for each block.""" index = [] for block_type, defn in self.blocks.items(): # Create a rich text representation for searching text = f"{block_type} {defn.get('category', '')} {defn.get('description', '')}" # Add property names to allow searching by property (e.g. "SRC" or "URL") text += " " + " ".join(defn.get('required_props', []) + defn.get('optional_props', [])) index.append({ 'id': block_type, 'text': text.lower(), 'data': defn }) return index def retrieve(self, query: str, top_k: int = 15) -> List[str]: """ Find top_k relevant blocks based on query. Uses simple keyword overlap and scoring for prototype. """ query_terms = set(query.lower().split()) scores = [] for item in self.index: score = 0 # 1. Exact phrase match in description if query.lower() in item['text']: score += 5 # 2. Term overlap item_terms = set(item['text'].split()) overlap = query_terms.intersection(item_terms) score += len(overlap) * 2 # 3. Fuzzy match for block type name for term in query_terms: if term in item['id']: score += 3 # Boost specific implementation details if "style" in query.lower() and "attr" in item['id']: score += 2 if score > 0: scores.append((score, item['id'])) # Always include essential root blocks if score is low but query implies structure # (For prototype, we just sort by score) scores.sort(key=lambda x: x[0], reverse=True) # Return top K IDs return [s[1] for s in scores[:top_k]] class RAGContextManager: """Manages conversation history and current workspace state.""" def __init__(self): self.history = [] self.current_workspace = {"blocks": [], "connections": []} def add_turn(self, user_input: str, ai_response: str): self.history.append({"user": user_input, "ai": ai_response}) def update_workspace(self, new_workspace: Dict): # In a real app, we would merge. For prototype, we replace. self.current_workspace = new_workspace def get_system_prompt_context(self) -> str: """Format current state for the AI""" if not self.current_workspace["blocks"]: return "Current Workspace: (Empty)" # summarize for token efficiency block_types = [b["type"] for b in self.current_workspace["blocks"]] return f"Current Workspace has {len(block_types)} blocks: {', '.join(block_types[:10])}..." class RAGAgent: def __init__(self): self.librarian = RAGLibrarian(BLOCKLY_SCHEMA) self.context = RAGContextManager() self.hf_token = os.getenv("HUGGINGFACE_TOKEN") self.client = InferenceClient(token=self.hf_token) # Use a smart model self.model = "Qwen/Qwen2.5-Coder-32B-Instruct" def construct_system_prompt(self, relevant_block_ids: List[str]) -> str: """Builds a dynamic system prompt with ONLY relevant blocks.""" # 1. Fetch definitions for retrieved blocks block_defs = [] for bid in relevant_block_ids: if bid in self.librarian.blocks: b = self.librarian.blocks[bid] # Format strictly for the LLM props = ", ".join(b.get('required_props', []) + b.get('optional_props', [])) block_defs.append(f"- {bid} (Category: {b['category']}): Props=[{props}]") # 2. Add essential infrastructure blocks (always needed) # e.g. root, body, text_content, attributes essentials = ['basic_html_root', 'basic_container', 'text_content', 'html_attr_class', 'text_value'] for eid in essentials: if eid not in relevant_block_ids and eid in self.librarian.blocks: b = self.librarian.blocks[eid] props = ", ".join(b.get('required_props', []) + b.get('optional_props', [])) block_defs.append(f"- {eid} (Category: {b['category']}): Props=[{props}]") blocks_str = "\n".join(block_defs) return f"""You are WebCraft AI, a RAG-powered agent. MISSING BLOCKS? The user might ask for something you don't have. finding the best match from the list below. AVAILABLE BLOCKS (Dynamically Retrieved): {blocks_str} RULES: 1. Output valid JSON with 'blocks' and 'connections'. 2. Use ONLY the blocks listed above. 3. If styling is asked, use `html_attr_class` with Tailwind CSS values connected to the target block. E.g. `basic_button` -> `ATTRS` -> `html_attr_class` -> `VALUE` -> `text_value`. CURRENT CONTEXT: {self.context.get_system_prompt_context()} """ def chat(self, user_input: str): print(f"\n🔍 [Librarian] Searching for blocks related to: '{user_input}'...") # 1. Retrieve relevant_ids = self.librarian.retrieve(user_input) print(f"📚 [Librarian] Found {len(relevant_ids)} relevant blocks: {relevant_ids}") # 2. Construct Prompt system_prompt = self.construct_system_prompt(relevant_ids) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_input} ] print("\n🤖 [Agent] Thinking...") try: response = self.client.chat_completion( messages=messages, model=self.model, max_tokens=2048, temperature=0.1 ) content = response.choices[0].message.content # Simple parsing for prototype display print("\n✨ [Agent] Response:") print(content[:500] + "..." if len(content) > 500 else content) # Try to parse JSON to simulate workspace update try: # Basic cleanup json_str = content if "```json" in content: json_str = content.split("```json")[1].split("```")[0] elif "```" in content: json_str = content.split("```")[1].split("```")[0] data = json.loads(json_str.strip()) self.context.update_workspace(data) print(f"\n✅ Valid JSON generated! {len(data.get('blocks', []))} blocks created.") except: print("\n⚠️ Response was not valid JSON (expected for partial prototype).") except Exception as e: print(f"❌ Error: {e}") if __name__ == "__main__": agent = RAGAgent() print("WebCraft RAG Prototype Initialized.") print("Type 'exit' to quit.") while True: try: user_input = input("\nUse Prompt >> ") if user_input.lower() in ['exit', 'quit']: break agent.chat(user_input) except KeyboardInterrupt: break