synth-ai 0.2.0__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

synth_ai/environments/examples/enron/art_helpers/local_email_db.py

@@ -0,0 +1,280 @@
+import sqlite3
+import os
+import logging
+from datasets import load_dataset, Dataset, Features, Value, Sequence
+from tqdm import tqdm
+from datetime import datetime
+
+# Resolve paths relative to this file so it works regardless of the current working directory
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+# Database will live in "../data/enron_emails.db" relative to project root
+DEFAULT_DB_PATH = os.path.join(BASE_DIR, "..", "..", "data", "enron_emails.db")
+
+DEFAULT_REPO_ID = "corbt/enron-emails"
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+# --- Database Schema ---
+SQL_CREATE_TABLES = """
+DROP TABLE IF EXISTS recipients;
+DROP TABLE IF EXISTS emails_fts;
+DROP TABLE IF EXISTS emails;
+
+CREATE TABLE emails (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    message_id TEXT UNIQUE,
+    subject TEXT,
+    from_address TEXT,
+    date TEXT, -- Store as ISO 8601 string 'YYYY-MM-DD HH:MM:SS'
+    body TEXT,
+    file_name TEXT
+);
+
+CREATE TABLE recipients (
+    email_id INTEGER,
+    recipient_address TEXT,
+    recipient_type TEXT, -- 'to', 'cc', 'bcc'
+    FOREIGN KEY(email_id) REFERENCES emails(id) ON DELETE CASCADE
+);
+"""
+
+SQL_CREATE_INDEXES_TRIGGERS = """
+CREATE INDEX idx_emails_from ON emails(from_address);
+CREATE INDEX idx_emails_date ON emails(date);
+CREATE INDEX idx_emails_message_id ON emails(message_id);
+CREATE INDEX idx_recipients_address ON recipients(recipient_address);
+CREATE INDEX idx_recipients_type ON recipients(recipient_type);
+CREATE INDEX idx_recipients_email_id ON recipients(email_id);
+CREATE INDEX idx_recipients_address_email ON recipients(recipient_address, email_id);
+
+CREATE VIRTUAL TABLE emails_fts USING fts5(
+    subject,
+    body,
+    content='emails',
+    content_rowid='id'
+);
+
+CREATE TRIGGER emails_ai AFTER INSERT ON emails BEGIN
+    INSERT INTO emails_fts (rowid, subject, body)
+    VALUES (new.id, new.subject, new.body);
+END;
+
+CREATE TRIGGER emails_ad AFTER DELETE ON emails BEGIN
+    DELETE FROM emails_fts WHERE rowid=old.id;
+END;
+
+CREATE TRIGGER emails_au AFTER UPDATE ON emails BEGIN
+    UPDATE emails_fts SET subject=new.subject, body=new.body WHERE rowid=old.id;
+END;
+
+INSERT INTO emails_fts (rowid, subject, body) SELECT id, subject, body FROM emails;
+"""
+
+
+# --- Functions ---
+
+
+def download_dataset(repo_id: str) -> Dataset:
+    """Downloads the dataset from Hugging Face Hub."""
+    logging.info(f"Attempting to download dataset from Hugging Face Hub: {repo_id}")
+    expected_features = Features(
+        {
+            "message_id": Value("string"),
+            "subject": Value("string"),
+            "from": Value("string"),
+            "to": Sequence(Value("string")),
+            "cc": Sequence(Value("string")),
+            "bcc": Sequence(Value("string")),
+            "date": Value("timestamp[us]"),
+            "body": Value("string"),
+            "file_name": Value("string"),
+        }
+    )
+    dataset_obj = load_dataset(repo_id, features=expected_features, split="train")
+    # Basic type check remains useful
+    if not isinstance(dataset_obj, Dataset):
+        raise TypeError(f"Expected Dataset, got {type(dataset_obj)}")
+    logging.info(f"Successfully loaded dataset '{repo_id}' with {len(dataset_obj)} records.")
+    return dataset_obj
+
+
+def create_database(db_path: str):
+    """Creates the SQLite database and tables."""
+    logging.info(f"Creating SQLite database and tables at: {db_path}")
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    cursor.executescript(SQL_CREATE_TABLES)
+    conn.commit()
+    conn.close()
+    logging.info("Database tables created successfully.")
+
+
+def populate_database(db_path: str, dataset: Dataset):
+    """Populates the database with data from the Hugging Face dataset."""
+    logging.info(f"Populating database {db_path}...")
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    # --- Performance Pragmas ---
+    conn.execute("PRAGMA synchronous = OFF;")
+    conn.execute("PRAGMA journal_mode = MEMORY;")
+
+    record_count = 0
+    skipped_count = 0  # Keep track of skipped emails due to filters
+    duplicate_count = 0  # Keep track of skipped duplicate emails
+    processed_emails = set()  # Track (subject, body, from) tuples to dedupe
+
+    conn.execute("BEGIN TRANSACTION;")  # Single transaction for bulk insert
+
+    for email_data in tqdm(dataset, desc="Inserting emails"):
+        assert isinstance(email_data, dict)
+        message_id = email_data["message_id"]
+        subject = email_data["subject"]
+        from_address = email_data["from"]
+        date_obj: datetime = email_data["date"]
+        body = email_data["body"]
+        file_name = email_data["file_name"]
+        to_list_raw = email_data["to"]
+        cc_list_raw = email_data["cc"]
+        bcc_list_raw = email_data["bcc"]
+
+        # --- Data Cleaning and Filtering ---
+        date_str = date_obj.strftime("%Y-%m-%d %H:%M:%S")
+        to_list = [str(addr) for addr in to_list_raw if addr]
+        cc_list = [str(addr) for addr in cc_list_raw if addr]
+        bcc_list = [str(addr) for addr in bcc_list_raw if addr]
+
+        # Check body length
+        if len(body) > 5000:
+            logging.debug(f"Skipping email {message_id}: Body length > 5000 characters.")
+            skipped_count += 1
+            continue
+
+        # Check total recipients
+        total_recipients = len(to_list) + len(cc_list) + len(bcc_list)
+        if total_recipients > 30:
+            logging.debug(
+                f"Skipping email {message_id}: Total recipients ({total_recipients}) > 30."
+            )
+            skipped_count += 1
+            continue
+        # --- End Filtering ---
+
+        # --- Deduplication Check ---
+        email_key = (subject, body, from_address)
+        if email_key in processed_emails:
+            logging.debug(
+                f"Skipping duplicate email (Subject: {subject[:50]}..., From: {from_address})"
+            )
+            duplicate_count += 1
+            continue
+        else:
+            processed_emails.add(email_key)
+        # --- End Deduplication ---
+
+        cursor.execute(
+            """
+            INSERT INTO emails (message_id, subject, from_address, date, body, file_name)
+            VALUES (?, ?, ?, ?, ?, ?)
+            """,
+            (message_id, subject, from_address, date_str, body, file_name),
+        )
+        email_pk_id = cursor.lastrowid
+
+        recipient_data = []
+        for addr in to_list:
+            recipient_data.append((email_pk_id, addr, "to"))
+        for addr in cc_list:
+            recipient_data.append((email_pk_id, addr, "cc"))
+        for addr in bcc_list:
+            recipient_data.append((email_pk_id, addr, "bcc"))
+
+        if recipient_data:
+            cursor.executemany(
+                """
+                INSERT INTO recipients (email_id, recipient_address, recipient_type)
+                VALUES (?, ?, ?)
+                """,
+                recipient_data,
+            )
+        record_count += 1
+
+    conn.commit()
+    conn.close()
+    logging.info(f"Successfully inserted {record_count} email records.")
+    if skipped_count > 0:
+        logging.info(f"Skipped {skipped_count} email records due to length or recipient limits.")
+    if duplicate_count > 0:
+        logging.info(
+            f"Skipped {duplicate_count} duplicate email records (based on subject, body, from)."
+        )
+
+
+def create_indexes_and_triggers(db_path: str):
+    """Creates indexes and triggers on the populated database."""
+    logging.info(f"Creating indexes and triggers for database: {db_path}...")
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    cursor.executescript(SQL_CREATE_INDEXES_TRIGGERS)
+    conn.commit()
+    conn.close()
+    logging.info("Indexes and triggers created successfully.")
+
+
+def generate_database(overwrite: bool = False):
+    """
+    Generates the SQLite database from the specified Hugging Face dataset.
+    Simplified version without extensive error handling.
+
+    Args:
+        repo_id: The Hugging Face repository ID for the dataset.
+        db_path: The path where the SQLite database file should be created.
+        overwrite: If True, any existing database file at db_path will be removed.
+    """
+    logging.info(
+        f"Starting database generation for repo '{DEFAULT_REPO_ID}' at '{DEFAULT_DB_PATH}'"
+    )
+    logging.info(f"Overwrite existing database: {overwrite}")
+
+    db_dir = os.path.dirname(DEFAULT_DB_PATH)
+    if db_dir and not os.path.exists(db_dir):
+        logging.info(f"Creating data directory: {db_dir}")
+        os.makedirs(db_dir)
+
+    if overwrite and os.path.exists(DEFAULT_DB_PATH):
+        logging.info(f"Removing existing database file: {DEFAULT_DB_PATH}")
+        os.remove(DEFAULT_DB_PATH)
+
+    if not os.path.exists(DEFAULT_DB_PATH):
+        dataset = download_dataset(DEFAULT_REPO_ID)
+        create_database(DEFAULT_DB_PATH)
+        populate_database(DEFAULT_DB_PATH, dataset)
+        create_indexes_and_triggers(DEFAULT_DB_PATH)
+
+        # ---- new: add unique index post-creation if not already handled by SQL_CREATE_INDEXES_TRIGGERS ---
+        # This ensures the index exists even if SQL_CREATE_INDEXES_TRIGGERS was modified
+        # or if we want to be absolutely certain this specific index is applied.
+        conn = sqlite3.connect(DEFAULT_DB_PATH)
+        cur = conn.cursor()
+        logging.info("Ensuring UNIQUE index on emails.message_id exists...")
+        cur.executescript(
+            """
+            CREATE UNIQUE INDEX IF NOT EXISTS idx_emails_message_id
+                ON emails(message_id);
+            """
+        )
+        conn.commit()
+        conn.close()
+        logging.info("UNIQUE index on emails.message_id verified/created.")
+        # ---- end new section ----
+    else:
+        logging.info(
+            f"Database already exists at {DEFAULT_DB_PATH}. Set overwrite=True to regenerate."
+        )
+
+    logging.info("Database generation process complete.")
+
+
+if __name__ == "__main__":
+    generate_database(overwrite=True)

synth_ai/environments/examples/enron/art_helpers/types_enron.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel
+from typing import List, Optional
+
+
+class SyntheticQuery(BaseModel):
+    id: int
+    question: str
+    answer: str
+    message_ids: List[str]  # message_ids (strings) of referenced emails
+    how_realistic: float
+    inbox_address: str
+    query_date: str
+
+
+class Email(BaseModel):
+    message_id: str
+    date: str  # ISO 8601 string 'YYYY-MM-DD HH:MM:SS'
+    subject: Optional[str] = None
+    from_address: Optional[str] = None
+    to_addresses: List[str] = []  # Populated from recipients table
+    cc_addresses: List[str] = []  # Populated from recipients table
+    bcc_addresses: List[str] = []  # Populated from recipients table
+    body: Optional[str] = None
+    file_name: Optional[str] = None

synth_ai/environments/examples/enron/engine.py

@@ -0,0 +1,291 @@
+# engine.py
+from __future__ import annotations
+from dataclasses import dataclass, asdict
+from typing import Any, Dict, Tuple, Optional, List
+from pydantic import BaseModel
+from pathlib import Path
+
+from synth_ai.environments.examples.enron.art_helpers.types_enron import Email
+from synth_ai.environments.examples.enron.art_helpers.email_search_tools import (
+    search_emails as helper_search_emails,
+    read_email as helper_read_email,
+    SearchResult,
+)
+
+# SQLite-backed helpers
+from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
+from synth_ai.environments.examples.enron.taskset import EnronTaskInstance
+from synth_ai.zyk import LM  # Import LM class
+
+from synth_ai.environments.environment.db.sqlite import SQLiteManager
+from synth_ai.environments.environment.rewards.core import RewardStack, RewardComponent
+from synth_ai.environments.examples.enron.art_helpers.local_email_db import (
+    DEFAULT_DB_PATH,
+    generate_database,
+)
+
+# --------------------------------------------------------------------------- actions
+ACTION_SEARCH = "search"
+ACTION_READ = "read"
+ACTION_ANSWER = "answer"
+
+
+# --------------------------------------------------------------------------- snapshot
+@dataclass
+class EnronEngineSnapshot(StatefulEngineSnapshot):
+    idx: int
+    answered: bool
+    total_reward: float
+    partial_rewards: List[float]
+
+
+# --------------------------------------------------------------------------- engine
+class EnronEngine(StatefulEngine):
+    """
+    Minimal state-machine around the corbt/enron_emails_sample_questions dataset.
+    Action is a tuple(kind, arg):
+
+        (ACTION_SEARCH,  query: str)      → returns {"search_results": [message_ids]}
+        (ACTION_READ,    message_id: str) → returns {"email_body": str}
+        (ACTION_ANSWER,  answer: str)     → rewards +1 / -1 and terminates
+    """
+
+    # ----------------------------- init / helpers
+    def __init__(self, task_instance: EnronTaskInstance):
+        # Use the provided TaskInstance snapshot for this episode
+        self.instance = task_instance
+        self.answered = False
+        self.total_reward = 0.0
+        self.idx = 0
+        # List to track each step's reward
+        self.rewards_history: List[float] = []
+
+        db_file_path = Path(DEFAULT_DB_PATH)
+        if not db_file_path.exists():
+            generate_database(overwrite=False)  # Ensure DB exists
+        self.sqlite_manager = SQLiteManager(db_path=db_file_path, read_only=True)
+
+        # RewardStack is an attribute of the engine; its calculations update private_state fields
+        self.reward_stack = RewardStack(
+            components=[
+                EnronAnswerCorrectnessComponent(),
+                EnronStepPenaltyComponent(penalty=-0.05),
+            ]
+        )
+        # This will hold the specific arguments/details of the current agent action
+        # for the reward components to inspect.
+        self._current_action_details_for_reward: Optional[Dict[str, Any]] = None
+
+    def _sample(self) -> Dict[str, Any]:
+        # Return the snapshot dict from the TaskInstance
+        return self.instance.initial_engine_snapshot
+
+    # ----------------------------- step / reset
+    async def _step_engine(
+        self, tool_output_payload: Optional[Dict[str, Any]]
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        r = await self._calculate_and_apply_reward()
+
+        # Determine termination: if an answer was attempted, task terminates.
+        # The 'answered' flag is set by answer_question_action.
+        term = self.answered
+
+        s = self._sample()
+        priv = {
+            "reward_last": r,
+            "total_reward": self.total_reward,
+            "terminated": term,
+            "truncated": False,
+            "gold_answer": s["answer"],
+        }
+
+        # Public state combines static elements with dynamic ones from tool_output_payload
+        pub = {
+            "question": s["question"],
+            "tools": [
+                "search_emails",
+                "read_email",
+                "answer_question",
+                "terminate",
+            ],  # Available tools
+            "already_answered": self.answered,
+            "query_date": s.get("query_date", "<unknown date>"),
+            "inbox_address": s.get("inbox_address", "<unknown_inbox>"),
+            # Default empty values, to be overwritten by tool_output_payload if present
+            "search_results": [],
+            "email": None,
+            **(tool_output_payload if tool_output_payload else {}),
+        }
+
+        return priv, pub
+
+    async def _reset_engine(
+        self, *, seed: Optional[int] = None
+    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """
+        Advance to the next Q-A pair and emit an initial observation **without**
+        issuing an empty-keyword DB search (which would raise).
+        """
+        # Reset answered status and total reward for this instance
+        self.answered = False
+        self.total_reward = 0.0
+        self.rewards_history = []
+        self._current_action_details_for_reward = None
+        # self.sqlite_manager.reset() # Enron DB is read-only; reset usually not needed unless switching DB files.
+
+        s = self._sample()
+        priv = {
+            "reward_last": 0.0,
+            "total_reward": 0.0,
+            "terminated": False,
+            "truncated": False,
+            "gold_answer": s["answer"],
+        }
+        pub = {
+            "question": s["question"],
+            "tools": ["search_emails", "read_email", "answer_question", "terminate"],
+            "already_answered": False,
+            "query_date": s.get("query_date", "<unknown date>"),
+            "inbox_address": s.get("inbox_address", "<unknown_inbox>"),
+            "search_results": [],
+            "email": None,
+        }
+        # No index advancement needed when using a single TaskInstance
+        return priv, pub
+
+    # ----------------------------- serialization helpers
+    async def _serialize_engine(self) -> EnronEngineSnapshot:
+        # Include partial rewards history in the snapshot
+        return EnronEngineSnapshot(
+            self.idx,
+            self.answered,
+            self.total_reward,
+            self.rewards_history,
+        )
+
+    @classmethod
+    async def _deserialize_engine(
+        cls, snap: EnronEngineSnapshot, task_instance: EnronTaskInstance
+    ) -> "EnronEngine":
+        eng = cls(task_instance)
+        eng.idx = snap.idx
+        eng.answered = snap.answered
+        eng.total_reward = snap.total_reward
+        eng.rewards_history = (
+            snap.partial_rewards
+        )  # Ensure this is correctly typed in Pydantic model if not List[float]
+        # Note: SQLiteManager is re-initialized in __init__ based on DEFAULT_DB_PATH.
+        # If the db path could change per instance/snapshot, that would need to be part of the snapshot.
+        return eng
+
+    def close_db(self):
+        self.sqlite_manager.close()
+
+    async def _calculate_and_apply_reward(self) -> float:
+        s = self._sample()
+        reward_context_state = {  # State snapshot for reward calculation
+            "question": s["question"],
+            "gold_answer": s["answer"],
+            **(
+                self._current_action_details_for_reward
+                if self._current_action_details_for_reward
+                else {}
+            ),
+        }
+
+        # The 'action' param for score can be the conceptual action type or detailed args
+        action_param_for_score = (
+            self._current_action_details_for_reward
+            if self._current_action_details_for_reward
+            else {}
+        )
+
+        reward = await self.reward_stack.step_reward(
+            state=reward_context_state, action=action_param_for_score
+        )
+
+        self.total_reward += reward
+        self.rewards_history.append(reward)
+        self._current_action_details_for_reward = None  # Reset after use
+        return reward
+
+    async def search_emails_action(self, search_args: Dict[str, Any]) -> List[Dict[str, Any]]:
+        res: List[SearchResult] = helper_search_emails(self.sqlite_manager, **search_args)
+        self._current_action_details_for_reward = {"type": "search", **search_args}
+        return [asdict(item) for item in res]
+
+    async def read_email_action(self, message_id: str) -> Optional[Dict[str, Any]]:
+        email: Optional[Email] = helper_read_email(self.sqlite_manager, message_id)
+        self._current_action_details_for_reward = {
+            "type": "read",
+            "message_id": message_id,
+        }
+        return email.dict() if email else None
+
+    async def answer_question_action(self, agent_answer: str) -> None:
+        # This method now primarily sets up state for reward calculation.
+        # The actual reward value and termination status are determined by _get_reward_and_update_state.
+        s = self._sample()
+        self._current_action_details_for_reward = {
+            "type": "answer",
+            "is_answer_action": True,  # Signal for reward component
+            "question": s["question"],
+            "gold_answer": s["answer"],
+            "agent_answer": agent_answer,
+        }
+        self.answered = True  # Mark as answered, termination decided by reward logic
+
+
+# ----------------------------- LLM Judge for answers
+async def determine_if_answer_is_correct(
+    question: str, gold_answer: str, agent_answer: str
+) -> bool:
+    # Instantiate LM for the judge
+    llm = LM(model_name="gpt-4.1-nano", formatting_model_name="gpt-4.1-nano", temperature=0.0)
+
+    system_prompt = (
+        "You will be given a question and two different answers to the question, "
+        "the correct answer and the answer given by an AI. Your job is to determine "
+        "if the answer given by the AI is correct."
+    )
+    user_message_content = (
+        f"Question: {question}\nCorrect answer: {gold_answer}\nAI answer: {agent_answer}"
+    )
+
+    class CorrectnessResponse(BaseModel):
+        correct: bool
+
+    # Use LM.respond_async
+    response = await llm.respond_async(
+        system_message=system_prompt,
+        user_message=user_message_content,
+        response_model=CorrectnessResponse,
+        # Caching is typically handled within the LM class or its underlying setup
+    )
+    return response.structured_output.correct
+
+
+# --- Placeholder Reward Components (ideally defined elsewhere and imported) ---
+# (These would typically live in a shared rewards components file or alongside the engine if very specific)
+class EnronAnswerCorrectnessComponent(RewardComponent):
+    async def score(self, state: Dict[str, Any], action: Any) -> float:
+        if state.get("is_answer_action") and state.get("agent_answer") is not None:
+            # determine_if_answer_is_correct should be part of the engine or accessible
+            # For now, assuming it's available in this scope.
+            correct = await determine_if_answer_is_correct(
+                state["question"], state["gold_answer"], state["agent_answer"]
+            )
+            return 1.0 if correct else -1.0
+        return 0.0
+
+
+class EnronStepPenaltyComponent(RewardComponent):
+    def __init__(self, penalty: float = -0.01):
+        self.penalty = penalty
+
+    async def score(self, state: Dict[str, Any], action: Any) -> float:
+        # Apply penalty for any action that isn't a final answer, or just every step.
+        # For simplicity, apply if not a "correct" answer action.
+        if not state.get("is_answer_action"):
+            return self.penalty
+        return 0.0