deploy_tmp/rag_utils.py

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)}"