""" 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 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, 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, "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] ) -> List[str]: """ Extract domain signature keywords for guardrail checking. """ # Start with analyzed keywords (highest priority) domain_keywords = prompt_analysis.get("domain_keywords", []) signature = list(domain_keywords) # Add domain and sub-domains domain = prompt_analysis.get("domain", "") if domain and domain not in signature: signature.insert(0, domain) for sub_domain in prompt_analysis.get("sub_domains", []): if sub_domain and sub_domain.lower() not in [s.lower() for s in signature]: signature.append(sub_domain) # Extract column headers from structured data for file_data in parsed_data: if file_data.get("data") and isinstance(file_data["data"], list): if file_data["data"] and isinstance(file_data["data"][0], dict): for key in file_data["data"][0].keys(): clean_key = key.lower().strip().replace("_", " ") if clean_key and clean_key not in [s.lower() for s in signature]: signature.append(clean_key) return signature[:80] # Limit for efficiency 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], system_prompt: str, prompt_analysis: Dict[str, Any], stats: Dict[str, Any] ): """Save agent artifacts to filesystem.""" 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, "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) 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)