""" Worker Agent 3 — Embedding Agent Task: Embed chunked documents into a vector store using selected model. Uses mock embeddings for determinism during training. """ from __future__ import annotations import hashlib import json import math from pathlib import Path from typing import Any, Optional from workers.base_worker import BaseWorker EMBEDDING_TASK_CONFIGS = { "easy_embedding": {"task_id": "easy_embedding", "difficulty": "easy", "step_budget": 8, "corpus_size": 50}, "medium_embedding": {"task_id": "medium_embedding", "difficulty": "medium", "step_budget": 10, "corpus_size": 100}, "hard_embedding": {"task_id": "hard_embedding", "difficulty": "hard", "step_budget": 12, "corpus_size": 200}, } VALID_MODELS = ["all-MiniLM-L6-v2", "all-mpnet-base-v2", "paraphrase-multilingual-MiniLM-L12-v2"] EMBEDDING_DIM = 384 STOP_WORDS = {"does", "is", "a", "the", "does", "support", "nexacrm", "faq", "what", "are", "you", "built", "on", "with", "how", "do", "i", "can", "an", "available", "on", "for", "in", "at", "rest"} def mock_embed(text: str) -> list[float]: """ Improved keyword-overlap embedding with stop-word filtering. Prioritizes unique terms like 'Slack', 'Pricing', or 'GDPR'. """ # Tokenization and filtering raw_words = text.lower().replace("?", "").replace("!", "").replace(".", "").replace(",", "").replace("\u2014", " ").split() words = [w for w in raw_words if w not in STOP_WORDS and len(w) > 1] if not words: # Fallback to all words if everything was a stopword words = [w for w in raw_words if len(w) > 1] if not words: return [0.0] * EMBEDDING_DIM vec = [0.0] * EMBEDDING_DIM for w in words: idx = int(hashlib.md5(w.encode()).hexdigest(), 16) % EMBEDDING_DIM # Keywords get higher weight vec[idx] += 1.0 norm = math.sqrt(sum(x * x for x in vec)) or 1.0 return [x / norm for x in vec] class EmbeddingEnv(BaseWorker): """ Worker Agent 3: Embedding Agent. Actions: - select_model: Choose embedding model - configure_batch_size: Set batch size for embedding - run_embedding: Embed all chunks - validate_coverage: Check coverage ratio - inspect_vectors: Sample 3 vectors - store_index: Persist index to memory - submit: Submit index and receive terminal reward """ VALID_ACTIONS = [ "select_model", "configure_batch_size", "run_embedding", "validate_coverage", "inspect_vectors", "store_index", "submit" ] def __init__(self) -> None: super().__init__(worker_id="worker_3", worker_name="Embedding Agent") self.selected_model: Optional[str] = None self.batch_size: int = 32 self.index: dict[str, list[float]] = {} # chunk_id -> vector self.chunks_embedded: int = 0 self.null_embedding_count: int = 0 self.coverage_ratio: float = 0.0 self.index_stored: bool = False self.embedding_done: bool = False self.task_config: dict = {} self.input_chunks: list[dict] = [] def reset(self, task_id: str, input_chunks: Optional[list[dict]] = None) -> dict: self._reset_episode_tracking() self.task_id = task_id self.task_config = EMBEDDING_TASK_CONFIGS.get(task_id, EMBEDDING_TASK_CONFIGS["easy_embedding"]) self.step_budget = self.task_config["step_budget"] self.step_budget_remaining = self.step_budget self.selected_model = None self.batch_size = 32 self.index = {} self.chunks_embedded = 0 self.null_embedding_count = 0 self.coverage_ratio = 0.0 self.index_stored = False self.embedding_done = False self.input_chunks = input_chunks or self._load_chunks(self.task_config["corpus_size"]) return self.state() def step(self, action_dict: dict) -> tuple[dict, float, bool, dict]: operation = action_dict.get("operation", "") parameters = action_dict.get("parameters", {}) if self.is_done: return self.state(), 0.0, True, {"error": "episode_already_done", "action": operation} if operation not in self.VALID_ACTIONS: self._record_governance_event("invalid_action", "medium", f"Unknown: {operation}") reward = self._clip_reward(0.0) self.step_count += 1 self.step_budget_remaining -= 1 self.total_reward += reward done = self._is_budget_exhausted() if done: self.is_done = True info = {"error": f"invalid_action:{operation}", "action": operation} self._record_action(operation, reward, info) return self.state(), reward, done, info reward, info = self._dispatch(operation, parameters) reward = self._clip_reward(reward) self.step_count += 1 self.step_budget_remaining -= 1 self.total_reward += reward done = self.is_done or self._is_budget_exhausted() if done: self.is_done = True self._record_action(operation, reward, info) return self.state(), reward, done, info def state(self) -> dict: base = self._get_base_state() base.update({ "selected_model": self.selected_model, "batch_size": self.batch_size, "chunks_embedded": self.chunks_embedded, "null_embedding_count": self.null_embedding_count, "coverage_ratio": round(self.coverage_ratio, 4), "index_stored": self.index_stored, "embedding_done": self.embedding_done, "input_chunk_count": len(self.input_chunks), }) return base def generate_run_report(self) -> dict: report = self._get_base_report() report.update({ "selected_model": self.selected_model, "chunks_embedded": self.chunks_embedded, "null_embedding_count": self.null_embedding_count, "coverage_ratio": self.coverage_ratio, "index_stored": self.index_stored, "final_score": self._compute_final_score(), }) return report def evaluate_run(self) -> dict: epsilon = 1e-6 score = min(max(self._compute_final_score(), epsilon), 1.0 - epsilon) gates = { "model_selected": self.selected_model is not None, "embedding_executed": self.embedding_done, "index_stored": self.index_stored, "no_null_embeddings": self.null_embedding_count == 0, "submitted": self.submitted, } return {"approved": all(gates.values()) and score >= 0.6, "gates": gates, "composite_score": score} def _dispatch(self, operation: str, parameters: dict) -> tuple[float, dict]: if operation == "select_model": return self._select_model(parameters) elif operation == "configure_batch_size": return self._configure_batch_size(parameters) elif operation == "run_embedding": return self._run_embedding() elif operation == "validate_coverage": return self._validate_coverage() elif operation == "inspect_vectors": return self._inspect_vectors() elif operation == "store_index": return self._store_index() elif operation == "submit": return self._submit() return 0.0, {"error": "unknown_operation"} def _select_model(self, params: dict) -> tuple[float, dict]: model = params.get("model_name", "") if model not in VALID_MODELS: return 0.1, {"error": f"invalid_model:{model}", "action": "select_model"} self.selected_model = model return 0.4, {"error": None, "action": "select_model", "model": model} def _configure_batch_size(self, params: dict) -> tuple[float, dict]: try: bs = int(params.get("batch_size", 32)) self.batch_size = max(1, min(bs, 256)) return 0.2, {"error": None, "action": "configure_batch_size", "batch_size": self.batch_size} except (ValueError, TypeError): return 0.0, {"error": "invalid_batch_size", "action": "configure_batch_size"} def _run_embedding(self) -> tuple[float, dict]: if not self.selected_model: self._record_governance_event("no_model", "high", "run_embedding called without model selected") return 0.05, {"error": "model_not_selected", "action": "run_embedding"} self.index = {} self.null_embedding_count = 0 for chunk in self.input_chunks: cid = chunk.get("chunk_id", f"chunk_{len(self.index):03d}") text = chunk.get("text", "") if not text.strip(): self.null_embedding_count += 1 continue self.index[cid] = mock_embed(text) self.chunks_embedded = len(self.index) total = len(self.input_chunks) self.coverage_ratio = self.chunks_embedded / max(total, 1) self.embedding_done = True reward = 0.4 * self.coverage_ratio + 0.3 * (1.0 if self.null_embedding_count == 0 else 0.3) + 0.2 * (self.step_budget_remaining / max(self.step_budget, 1)) return reward, {"error": None, "action": "run_embedding", "chunks_embedded": self.chunks_embedded} def _validate_coverage(self) -> tuple[float, dict]: if not self.embedding_done: return 0.1, {"error": "embedding_not_done", "action": "validate_coverage"} report = {"coverage_ratio": round(self.coverage_ratio, 4), "null_count": self.null_embedding_count, "total_chunks": len(self.input_chunks)} reward = 0.5 if self.coverage_ratio >= 0.95 else 0.2 return reward, {"error": None, "action": "validate_coverage", "report": report} def _inspect_vectors(self) -> tuple[float, dict]: sample = {k: v[:5] for k, v in list(self.index.items())[:3]} return 0.15, {"error": None, "action": "inspect_vectors", "sample": sample} def _store_index(self) -> tuple[float, dict]: if not self.embedding_done: return 0.05, {"error": "embedding_not_done", "action": "store_index"} self.index_stored = True return 0.4, {"error": None, "action": "store_index", "index_size": len(self.index)} def _submit(self) -> tuple[float, dict]: if not self.index_stored: self._record_governance_event("premature_submit", "high", "submit called before store_index") return 0.05, {"error": "index_not_stored", "action": "submit"} self.submitted = True self.is_done = True final_score = self._compute_final_score() return min(final_score + 0.15, 0.99), {"error": None, "action": "submit", "final_score": final_score} def _compute_final_score(self) -> float: if not self.embedding_done: return 0.001 score = ( 0.4 * self.coverage_ratio + 0.3 * (1.0 if self.null_embedding_count == 0 else 0.0) + 0.2 * (1.0 if self.index_stored else 0.0) + 0.1 * (self.step_budget_remaining / max(self.step_budget, 1)) ) return self._clip_reward(score) def _load_chunks(self, size: int) -> list[dict]: path = Path("data/nexacrm_corpus.json") if path.exists(): with open(path) as f: return json.load(f)[:size] return [{"chunk_id": f"chunk_{i:03d}", "text": f"NexaCRM entry {i}"} for i in range(size)]