import difflib from typing import List, Dict, Any import json try: from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity SKLEARN_AVAILABLE = True except ImportError: SKLEARN_AVAILABLE = False class RAGLibrarian: """The 'Librarian' that finds relevant blocks from the schema.""" def __init__(self, schema: Dict): self.blocks = schema.get('blocks', {}) self.index = self._build_index() if SKLEARN_AVAILABLE: # Better tokenization: split on underscores so "html_button" indexes as "html" and "button" self.vectorizer = TfidfVectorizer(stop_words='english', token_pattern=r'(?u)[a-zA-Z0-9]+') self.corpus = [item['text'] for item in self.index] if self.corpus: self.tfidf_matrix = self.vectorizer.fit_transform(self.corpus) 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} {block_type.replace('_', ' ')} {defn.get('category', '')} {defn.get('description', '')}" # Add property names and values context text += " " + " ".join(defn.get('required_props', []) + defn.get('optional_props', [])) # Add implicit semantic keywords based on block type if 'img' in block_type or 'image' in block_type: text += " picture photo logo visual graphic image" if 'style' in block_type or 'attr' in block_type: text += " color css design background padding margin font bold styling visual theme" if 'button' in block_type: text += " clickable action submit link btn button" if 'nav' in block_type: text += " header menu links topbar navigation nav" if 'section' in block_type or 'container' in block_type: text += " layout wrapper grouping box area part section div container hero" if 'card' in block_type: text += " panel box component widget card tile" index.append({ 'id': block_type, 'text': text.lower(), 'data': defn }) return index def retrieve(self, query: str, top_k: int = 20) -> List[str]: """ Find top_k relevant blocks based on query. Uses TF-IDF Cosine Similarity if available, otherwise falls back to simple overlap. """ if not SKLEARN_AVAILABLE or not self.corpus: return self._retrieve_fallback(query, top_k) # 1. Expand query slightly based on intent expanded_query = query.lower() if "style" in expanded_query or "color" in expanded_query: expanded_query += " attr style class" # 2. Vectorize user query query_vec = self.vectorizer.transform([expanded_query]) # 3. Compute cosine similarity against all blocks similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten() if similarities.max() == 0: return self._retrieve_fallback(query, top_k) # 4. Get top K indices top_indices = similarities.argsort()[-top_k:][::-1] # 5. Return Block IDs return [self.index[i]['id'] for i in top_indices if similarities[i] > 0] def _retrieve_fallback(self, query: str, top_k: int) -> List[str]: """Simple keyword overlap fallback if sklearn is not installed.""" query_terms = set(query.lower().split()) scores = [] for item in self.index: score = 0 if query.lower() in item['text']: score += 5 item_terms = set(item['text'].split()) score += len(query_terms.intersection(item_terms)) * 2 for term in query_terms: if term in item['id']: score += 3 if "style" in query.lower() and "attr" in item['id']: score += 2 if score > 0: scores.append((score, item['id'])) scores.sort(key=lambda x: x[0], reverse=True) return [s[1] for s in scores[:top_k]] class RAGContextManager: """Manages context and summarization for the RAG agent.""" @staticmethod def summarize_workspace(workspace_json: str) -> str: """ Convert a Blockly workspace JSON string into a structured tree hierarchy that gives the LLM Spatial Awareness of layouts. """ import json try: if not workspace_json: return "No existing workspace." data = json.loads(workspace_json) blocks = data.get("blocks", []) connections = data.get("connections", []) if not blocks: return "Workspace is empty." # Map blocks by ID for easy lookup block_map = {b['id']: b for b in blocks} # Build connection maps # child_to_parent maps ChildID -> ParentID # parent_to_children maps ParentID -> List of ChildIDs child_to_parent = {} parent_to_children = {} # Connections come as: ["parentId.INPUT_NAME", "childId"] or ["prevId.NEXT", "nextId"] # To build a visual tree, we track purely hierarchical "Inside/Under" relationships. # Next block connections will be rendered at the same indentation level. parent_inputs = {} next_links = {} for conn in connections: source, target_id = conn if "." in source: source_id, input_name = source.split(".", 1) if input_name == "NEXT": next_links[source_id] = target_id else: parent_inputs.setdefault(source_id, []).append((input_name, target_id)) child_to_parent[target_id] = source_id # Find root nodes (blocks with no parent and no previous sibling pointing to them) roots = [] for b_id in block_map: if b_id not in child_to_parent and b_id not in next_links.values(): roots.append(b_id) summary = ["CURRENT WORKSPACE LAYOUT (HIERARCHY):", f"Total Blocks: {len(blocks)}\n"] def render_tree(node_id, depth=0): if node_id not in block_map: return block = block_map[node_id] indent = " " * depth # Format properties props = block.get('props', {}) visible_props = {k: v for k, v in props.items() if k not in ['x', 'y', 'id']} prop_str = f" Props: {visible_props}" if visible_props else "" # Output current node summary.append(f"{indent}- [{block['id']}] {block['type']}{prop_str}") # Render children (inputs) if node_id in parent_inputs: for input_name, child_id in parent_inputs[node_id]: summary.append(f"{indent} (Input: {input_name} ->)") render_tree(child_id, depth + 2) # Render NEXT sibling at same depth if node_id in next_links: render_tree(next_links[node_id], depth) for root_id in roots: render_tree(root_id, 0) return "\n".join(summary) except Exception as e: return f"Error summarizing workspace: {str(e)}"