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