""" MEXAR Core Engine - Knowledge Compilation Module Builds Vector embeddings from parsed data for semantic retrieval. """ import os import json import logging from typing import Dict, List, Any, Optional from pathlib import Path from utils.groq_client import get_groq_client, GroqClient from utils.domain_signature import build_domain_signature from fastembed import TextEmbedding from core.database import SessionLocal from models.agent import Agent from models.chunk import DocumentChunk # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class KnowledgeCompiler: """ Compiles knowledge from parsed data into Vector embeddings. Uses semantic similarity for retrieval-based reasoning. """ def __init__(self, groq_client: Optional[GroqClient] = None, data_dir: str = "data/agents"): """ Initialize the knowledge compiler. Args: groq_client: Optional pre-configured Groq client data_dir: Directory to store agent data """ self.client = groq_client or get_groq_client() self.data_dir = Path(data_dir) self.data_dir.mkdir(parents=True, exist_ok=True) # Compilation progress tracking self.progress = { "status": "idle", "percentage": 0, "current_step": "", "details": {} } # Initialize embedding model (384 dim default) try: # Force cache to /tmp for HF Spaces or use env var cache_dir = os.getenv("FASTEMBED_CACHE_PATH", "/tmp/.cache/fastembed") self.embedding_model = TextEmbedding( model_name="BAAI/bge-small-en-v1.5", cache_dir=cache_dir ) logger.info(f"FastEmbed model loaded (cache: {cache_dir})") except Exception as e: logger.warning(f"Failed to load embedding model: {e}") self.embedding_model = None def compile( self, agent_name: str, parsed_data: List[Dict[str, Any]], system_prompt: str, prompt_analysis: Dict[str, Any] ) -> Dict[str, Any]: """ Main compilation function. Args: agent_name: Name of the agent being created parsed_data: List of parsed file results from DataValidator system_prompt: User's system prompt prompt_analysis: Analysis from PromptAnalyzer Returns: Dict containing: - domain_signature: Keywords for domain matching - stats: Compilation statistics """ self._update_progress("starting", 0, "Initializing compilation...") try: # Step 1: Build text context (30%) self._update_progress("building_context", 10, "Building text context...") text_context = self._build_text_context(parsed_data) self._update_progress("building_context", 30, f"Text context built: {len(text_context):,} characters") # Step 2: Extract domain signature (50%) self._update_progress("extracting_signature", 35, "Extracting domain signature...") domain_signature = self._extract_domain_signature(parsed_data, prompt_analysis) self._update_progress("extracting_signature", 50, f"Domain signature: {len(domain_signature)} keywords") # Step 3: Calculate stats (60%) self._update_progress("calculating_stats", 55, "Calculating statistics...") stats = self._calculate_stats(text_context, parsed_data) # Step 4: Save metadata (70%) self._update_progress("saving", 65, "Saving agent metadata...") self._save_agent( agent_name=agent_name, text_context=text_context, domain_signature=domain_signature["combined"], domain_signature_weights=domain_signature["lexical_weights"], domain_entities=domain_signature["entities"], system_prompt=system_prompt, prompt_analysis=prompt_analysis, stats=stats ) # Step 5: Save to Vector DB (95%) if self.embedding_model: self._update_progress("saving_vector", 75, "Saving to Vector Store...") self._save_to_vector_db(agent_name, text_context) self._update_progress("complete", 100, "Compilation complete!") return { "domain_signature": domain_signature["combined"], "stats": stats, "agent_path": str(self.data_dir / agent_name) } except Exception as e: logger.error(f"Compilation failed: {e}") self._update_progress("error", self.progress["percentage"], f"Error: {str(e)}") raise def _update_progress(self, status: str, percentage: int, step: str, details: Dict = None): """Update compilation progress.""" self.progress = { "status": status, "percentage": percentage, "current_step": step, "details": details or {} } logger.info(f"[{percentage}%] {step}") def get_progress(self) -> Dict[str, Any]: """Get current compilation progress.""" return self.progress.copy() def _build_text_context(self, parsed_data: List[Dict[str, Any]]) -> str: """ Build text context from parsed data. Args: parsed_data: Parsed file data Returns: Formatted text context """ context_parts = [] for i, file_data in enumerate(parsed_data): file_name = file_data.get("file_name", file_data.get("source", f"Source_{i+1}")) file_format = file_data.get("format", file_data.get("type", "unknown")) context_parts.append(f"\n{'='*60}") context_parts.append(f"SOURCE: {file_name} ({file_format.upper()})") context_parts.append(f"{'='*60}\n") # Handle structured data (CSV, JSON) if file_data.get("data"): for j, entry in enumerate(file_data["data"]): if isinstance(entry, dict): entry_lines = [f"[Entry {j+1}]"] for key, value in entry.items(): if value is not None and str(value).strip(): entry_lines.append(f" {key}: {value}") context_parts.append("\n".join(entry_lines)) else: context_parts.append(f"[Entry {j+1}] {entry}") # Handle unstructured text (PDF, DOCX, TXT) elif file_data.get("text"): context_parts.append(file_data["text"]) # Handle content field elif file_data.get("content"): context_parts.append(file_data["content"]) # Handle records field elif file_data.get("records"): for j, record in enumerate(file_data["records"]): if record and record.strip(): context_parts.append(f"[Line {j+1}] {record}") text_context = "\n".join(context_parts) # Limit to prevent token overflow (approximately 128K tokens = 500K chars) max_chars = 500000 if len(text_context) > max_chars: logger.warning(f"Text context truncated from {len(text_context)} to {max_chars} characters") text_context = text_context[:max_chars] + "\n\n[CONTEXT TRUNCATED DUE TO SIZE LIMITS]" return text_context def _extract_domain_signature( self, parsed_data: List[Dict[str, Any]], prompt_analysis: Dict[str, Any] ) -> Dict: """ Build domain signature using TF-IDF (Eq. 1) + NER entities (Section III-A). Extracts per-file raw text to preserve per-document granularity needed for TF-IDF's document-frequency term to be meaningful. Returns the full signature dict from build_domain_signature() including 'combined', 'lexical_weights', and 'entities' keys. """ # Collect per-file text documents (NOT a single concatenated string) documents: List[str] = [] for file_data in parsed_data: text = ( file_data.get("text") or file_data.get("content") or "" ) # For structured data (CSV/JSON), build a text representation per file if not text and file_data.get("data") and isinstance(file_data["data"], list): rows = file_data["data"] text = " ".join( " ".join(str(v) for v in row.values() if v is not None) for row in rows if isinstance(row, dict) ) if text and text.strip(): documents.append(text) # Build TF-IDF + NER domain signature signature = build_domain_signature( documents=documents, tau_tf=0.0, tau_ent=2, top_n_lexical=100, ) logger.info( f"Domain signature: {len(signature['combined'])} combined terms " f"({len(signature['lexical'])} lexical + {len(signature['entities'])} NER entities)" ) return signature def _calculate_stats( self, text_context: str, parsed_data: List[Dict[str, Any]] ) -> Dict[str, Any]: """Calculate compilation statistics.""" return { "context_length": len(text_context), "context_tokens": len(text_context) // 4, # Rough estimate "source_files": len(parsed_data), "total_entries": sum( len(p.get("data", [])) or len(p.get("records", [])) for p in parsed_data ) } def _save_agent( self, agent_name: str, text_context: str, domain_signature: List[str], domain_signature_weights: Dict, domain_entities: List[str], system_prompt: str, prompt_analysis: Dict[str, Any], stats: Dict[str, Any] ): """Save agent artifacts to filesystem and Supabase DB.""" agent_dir = self.data_dir / agent_name agent_dir.mkdir(parents=True, exist_ok=True) # Save text context (for backup/debugging) with open(agent_dir / "context.txt", "w", encoding="utf-8") as f: f.write(text_context) # Save metadata metadata = { "agent_name": agent_name, "system_prompt": system_prompt, "prompt_analysis": prompt_analysis, "domain_signature": domain_signature, "domain_signature_weights": domain_signature_weights, "domain_entities": domain_entities, "stats": stats, "created_at": self._get_timestamp() } with open(agent_dir / "metadata.json", "w", encoding="utf-8") as f: json.dump(metadata, f, indent=2, ensure_ascii=False) # Persist new domain_signature_weights and domain_entities to Supabase DB try: with SessionLocal() as db: agent = db.query(Agent).filter(Agent.name == agent_name).first() if agent: agent.domain_signature = domain_signature agent.domain_signature_weights = domain_signature_weights agent.domain_entities = domain_entities db.commit() logger.info( f"Persisted domain_signature ({len(domain_signature)} terms), " f"domain_signature_weights, and domain_entities for agent '{agent_name}'." ) except Exception as e: logger.warning(f"Could not persist domain signature to DB: {e}") logger.info(f"Agent saved to: {agent_dir}") def _get_timestamp(self) -> str: """Get current timestamp.""" from datetime import datetime return datetime.now().isoformat() def load_agent(self, agent_name: str) -> Dict[str, Any]: """ Load a previously compiled agent. Args: agent_name: Name of the agent to load Returns: Dict with agent artifacts """ agent_dir = self.data_dir / agent_name if not agent_dir.exists(): raise FileNotFoundError(f"Agent '{agent_name}' not found") # Load metadata with open(agent_dir / "metadata.json", "r", encoding="utf-8") as f: metadata = json.load(f) return { "metadata": metadata, "domain_signature": metadata.get("domain_signature", []), "system_prompt": metadata.get("system_prompt", ""), "prompt_analysis": metadata.get("prompt_analysis", {}) } def _save_to_vector_db(self, agent_name: str, context: str): """Chunk and save context to vector database.""" try: chunks = self._chunk_text(context) if not chunks: logger.warning(f"No chunks generated for {agent_name}") return logger.info(f"Generating embeddings for {len(chunks)} chunks...") # Generate embeddings with error handling try: embeddings = list(self.embedding_model.embed(chunks)) logger.info(f"Successfully generated {len(embeddings)} embeddings") except Exception as embed_error: logger.error(f"Embedding generation failed: {embed_error}") # Don't fail the entire compilation if embeddings fail return with SessionLocal() as db: agent = db.query(Agent).filter(Agent.name == agent_name).first() if not agent: logger.error(f"Agent {agent_name} not found in DB") return # Clear old chunks try: deleted_count = db.query(DocumentChunk).filter(DocumentChunk.agent_id == agent.id).delete() logger.info(f"Deleted {deleted_count} old chunks for agent {agent_name}") except Exception as delete_error: logger.warning(f"Failed to delete old chunks: {delete_error}") # Continue anyway # Insert new chunks try: new_chunks = [ DocumentChunk( agent_id=agent.id, content=chunk, embedding=embedding.tolist(), source="context" ) for chunk, embedding in zip(chunks, embeddings) ] db.add_all(new_chunks) # Update agent's chunk_count agent.chunk_count = len(new_chunks) db.commit() logger.info(f"Saved {len(new_chunks)} chunks to vector store for {agent_name}") except Exception as insert_error: logger.error(f"Failed to insert chunks: {insert_error}") db.rollback() raise except Exception as e: logger.error(f"Vector save failed: {e}", exc_info=True) # Don't raise - allow compilation to continue even if vector save fails def _chunk_text(self, text: str, chunk_size: int = 1000, overlap: int = 100) -> List[str]: """Simple text chunker.""" chunks = [] if not text: return [] start = 0 while start < len(text): end = min(start + chunk_size, len(text)) chunks.append(text[start:end]) if end == len(text): break start += (chunk_size - overlap) return chunks def list_agents(self) -> List[Dict[str, Any]]: """List all compiled agents.""" agents = [] for agent_dir in self.data_dir.iterdir(): if agent_dir.is_dir(): metadata_path = agent_dir / "metadata.json" if metadata_path.exists(): with open(metadata_path, "r", encoding="utf-8") as f: metadata = json.load(f) agents.append({ "name": agent_dir.name, "domain": metadata.get("prompt_analysis", {}).get("domain", "unknown"), "created_at": metadata.get("created_at"), "stats": metadata.get("stats", {}) }) return agents # Factory function def create_knowledge_compiler(data_dir: str = "data/agents") -> KnowledgeCompiler: """Create a new KnowledgeCompiler instance.""" return KnowledgeCompiler(data_dir=data_dir)