synth-ai 0.2.14__py3-none-any.whl → 0.4.4__py3-none-any.whl

synth_ai/environments/examples/enron/engine.py

@@ -1,300 +0,0 @@
-# engine.py
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
-
-from pydantic import BaseModel
-from synth_ai.environments.environment.db.sqlite import SQLiteManager
-from synth_ai.environments.environment.rewards.core import RewardComponent, RewardStack
-from synth_ai.environments.examples.enron.art_helpers.email_search_tools import (
-    SearchResult,
-)
-from synth_ai.environments.examples.enron.art_helpers.email_search_tools import (
-    read_email as helper_read_email,
-)
-from synth_ai.environments.examples.enron.art_helpers.email_search_tools import (
-    search_emails as helper_search_emails,
-)
-from synth_ai.environments.examples.enron.art_helpers.local_email_db import (
-    DEFAULT_DB_PATH,
-    generate_database,
-)
-from synth_ai.environments.examples.enron.art_helpers.types_enron import Email
-from synth_ai.environments.examples.enron.taskset import EnronTaskInstance
-
-# SQLite-backed helpers
-from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
-try:  # pragma: no cover - optional dependency
-    from synth_ai.zyk import LM  # type: ignore
-except ImportError:  # pragma: no cover - fallback when LM unavailable
-    LM = None
-
-# --------------------------------------------------------------------------- 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:
-    if LM is None:
-        return gold_answer.strip().lower() == agent_answer.strip().lower()
-
-    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

synth_ai/environments/examples/enron/environment.py

@@ -1,234 +0,0 @@
-# environment.py
-from __future__ import annotations
-
-from typing import Any, Dict, List, Optional, Union
-
-from pydantic import BaseModel, Field
-from synth_ai.environments.environment.shared_engine import (
-    GetObservationCallable,
-    InternalObservation,
-)
-from synth_ai.environments.environment.tools import (
-    AbstractTool,
-    TOOL_REGISTRY,
-    EnvToolCall,
-    ToolResult,
-    register_tool,
-)
-from synth_ai.environments.examples.enron.engine import (
-    ACTION_ANSWER,
-    ACTION_READ,
-    ACTION_SEARCH,
-    EnronEngine,
-)
-from synth_ai.environments.examples.enron.taskset import EnronTaskInstance
-from synth_ai.environments.stateful.core import StatefulEnvironment
-
-
-# -------- pydantic schemas (used by agent / LLM function calls)
-class SearchEmailsArgs(BaseModel):
-    inbox: str = Field(..., description="Email address performing the search (used by tool logic)")
-    keywords: List[str] = Field(..., description="Keywords to AND-search for")
-    from_addr: Optional[str] = None
-    to_addr: Optional[str] = None
-    sent_after: Optional[str] = None
-    sent_before: Optional[str] = None
-    max_results: int = Field(10, le=10)
-
-
-class ReadEmailArgs(BaseModel):
-    message_id: str
-
-
-class AnswerQuestionArgs(BaseModel):
-    answer: str
-
-
-# --------------------------------------------------------------------------- tool wrappers
-class SearchEmails(EnvToolCall):
-    def __init__(self, **kwargs):
-        self.action = (ACTION_SEARCH, kwargs)
-
-
-class ReadEmail(EnvToolCall):
-    def __init__(self, message_id: str):
-        self.action = (ACTION_READ, message_id)
-
-
-class AnswerQuestion(EnvToolCall):
-    def __init__(self, answer: str):
-        self.action = (ACTION_ANSWER, answer)
-
-
-# -- terminate wrapper (maps to an empty-answer ACTION_ANSWER) --------------
-class Terminate(EnvToolCall):
-    def __init__(self):
-        self.action = (ACTION_ANSWER, "")
-
-
-class TerminateArgs(BaseModel):
-    pass
-
-
-class SearchEmailsTool(AbstractTool):
-    name = "search_emails"
-    call_schema = SearchEmailsArgs
-
-    def __init__(self, engine: EnronEngine):
-        self.engine = engine
-
-    async def __call__(self, call: EnvToolCall) -> ToolResult:
-        try:
-            args = self.call_schema.model_validate(call.args or {})
-            results = await self.engine.search_emails_action(args.model_dump())
-            return ToolResult(ok=True, payload={"search_results": results})
-        except Exception as exc:  # pragma: no cover - runtime safety
-            return ToolResult(ok=False, error=str(exc))
-
-
-class ReadEmailTool(AbstractTool):
-    name = "read_email"
-    call_schema = ReadEmailArgs
-
-    def __init__(self, engine: EnronEngine):
-        self.engine = engine
-
-    async def __call__(self, call: EnvToolCall) -> ToolResult:
-        try:
-            args = self.call_schema.model_validate(call.args or {})
-            email = await self.engine.read_email_action(args.message_id)
-            return ToolResult(ok=True, payload={"email": email})
-        except Exception as exc:  # pragma: no cover
-            return ToolResult(ok=False, error=str(exc))
-
-
-class AnswerQuestionTool(AbstractTool):
-    name = "answer_question"
-    call_schema = AnswerQuestionArgs
-
-    def __init__(self, engine: EnronEngine):
-        self.engine = engine
-
-    async def __call__(self, call: EnvToolCall) -> ToolResult:
-        try:
-            args = self.call_schema.model_validate(call.args or {})
-            await self.engine.answer_question_action(args.answer)
-            return ToolResult(ok=True, payload={"status": "answer_recorded"})
-        except Exception as exc:  # pragma: no cover
-            return ToolResult(ok=False, error=str(exc))
-
-
-class TerminateTool(AbstractTool):
-    name = "terminate"
-    call_schema = TerminateArgs
-
-    def __init__(self, engine: EnronEngine):
-        self.engine = engine
-
-    async def __call__(self, call: EnvToolCall) -> ToolResult:
-        try:
-            await self.engine.answer_question_action("")
-            return ToolResult(ok=True, payload={"status": "terminated"})
-        except Exception as exc:  # pragma: no cover
-            return ToolResult(ok=False, error=str(exc))
-
-
-# -------- observation callable (optional for formatted observations)
-class SynthEnronObservationCallable(GetObservationCallable):
-    async def get_observation(
-        self, pub: Dict[str, Any], priv: Dict[str, Any]
-    ) -> InternalObservation:
-        """Format observation as a human-readable string."""
-        q = pub.get("question")
-        rwd = priv.get("reward_last")
-        return f"Q: {q}\nTools: {pub.get('tools')}\nAnswered: {pub.get('already_answered')}\nSearch Res: {len(pub.get('search_results', []))} items\nEmail Loaded: {pub.get('email') is not None}\nTool Error: {pub.get('tool_error')}\nReward Δ: {rwd}"
-
-
-# --------------------------------------------------------------------------- environment
-class EnronEnvironment(StatefulEnvironment):
-    def __init__(
-        self,
-        task_instance: EnronTaskInstance,
-        custom_obs: Optional[GetObservationCallable] = None,
-    ):
-        self.engine = EnronEngine(task_instance)
-        self.custom_obs = custom_obs or SynthEnronObservationCallable()
-        self.name = "Enron-QA-Env"
-
-        # Store tool instances on self for reliable access
-        self._tools_instances = {
-            "search_emails": SearchEmailsTool(self.engine),
-            "read_email": ReadEmailTool(self.engine),
-            "answer_question": AnswerQuestionTool(self.engine),
-            "terminate": TerminateTool(self.engine),
-        }
-        for tool_name, tool_instance in self._tools_instances.items():
-            if tool_name not in TOOL_REGISTRY:
-                register_tool(tool_instance)
-            elif TOOL_REGISTRY[tool_name].engine is not self.engine:
-                register_tool(tool_instance)
-
-    async def initialize(self) -> InternalObservation:
-        priv, pub = await self.engine._reset_engine()
-        return await self._obs(priv, pub)
-
-    async def step(
-        self,
-        calls: Union[EnvToolCall, List[EnvToolCall], List[List[EnvToolCall]]],
-    ) -> InternalObservation:
-        # normalise → always [[EnvToolCall]]
-        if isinstance(calls, EnvToolCall):
-            calls = [[calls]]
-        elif calls and isinstance(calls[0], EnvToolCall):
-            calls = [calls]
-
-        if not isinstance(calls[0][0], EnvToolCall):
-            raise TypeError(f"Processed call is not EnvToolCall: {type(calls[0][0])}")
-
-        tool_name = calls[0][0].tool
-        tool_to_execute = self._tools_instances.get(tool_name)
-
-        if not tool_to_execute:
-            tool_to_execute = TOOL_REGISTRY.get(tool_name)
-            if not tool_to_execute:
-                raise ValueError(f"Tool '{tool_name}' not found.")
-
-        tool_result: ToolResult = await tool_to_execute(calls[0][0])
-
-        public_payload_for_engine = (
-            tool_result.payload if tool_result.ok and tool_result.payload else {}
-        )
-        if not tool_result.ok:
-            public_payload_for_engine["tool_error"] = tool_result.error
-
-        priv, pub = await self.engine._step_engine(public_payload_for_engine)
-        return await self._obs(priv, pub)
-
-    async def terminate(self) -> InternalObservation:
-        self.engine.close_db()
-        priv_state_on_terminate = {
-            "reward_last": 0,
-            "total_reward": self.engine.total_reward,
-            "terminated": True,
-            "truncated": False,
-            "gold_answer": self.engine._sample()["answer"],
-        }
-        pub_state_on_terminate = {
-            "question": self.engine._sample()["question"],
-            "tools": [],
-            "already_answered": self.engine.answered,
-            "status": "terminated_by_env",
-        }
-        return await self._obs(priv_state_on_terminate, pub_state_on_terminate)
-
-    async def checkpoint(self) -> InternalObservation:
-        snapshot = await self.engine._serialize_engine()
-        return {
-            "engine_snapshot": snapshot.model_dump(),
-            "message": "Checkpoint created",
-        }
-
-    async def _obs(self, priv: Dict[str, Any], pub: Dict[str, Any]):
-        if self.custom_obs:
-            return await self.custom_obs.get_observation(pub, priv)
-        return {**pub, **priv}

synth_ai/environments/examples/enron/taskset.py

@@ -1,112 +0,0 @@
-# taskset.py
-from __future__ import annotations
-
-import asyncio
-import os
-from dataclasses import asdict, dataclass
-from uuid import uuid4
-
-from datasets import load_dataset
-from synth_ai.environments.tasks.core import (
-    Impetus,
-    Intent,
-    SplitInfo,
-    Task,
-    TaskInstance,
-    TaskInstanceMetadata,
-    TaskInstanceSet,
-)
-
-enron_task = Task(
-    global_premises="Answer factual questions by reading Enron e-mails",
-    global_constraints="",
-    global_objectives="Provide the correct answer; minimise queries",
-    shared_env_params={},
-)
-
-
-# --------------------------------------------------------------------------- metadata
-@dataclass
-class EnronTaskInstanceMetadata(TaskInstanceMetadata):
-    split: str
-    email_count: int
-    message_ids: list[str]
-
-
-@dataclass
-class EnronTaskInstance(TaskInstance):
-    async def serialize(self):
-        data = asdict(self)
-        if isinstance(data.get("id"), uuid4().__class__):
-            data["id"] = str(data["id"])
-        return data
-
-    @classmethod
-    async def deserialize(cls, data: dict) -> "EnronTaskInstance":
-        return cls(**data)
-
-
-# --------------------------------------------------------------------------- task-set builder
-# Use a local dataset cache under examples/enron/dataset
-CACHE_DIR = os.path.join(os.path.dirname(__file__), "dataset")
-os.makedirs(CACHE_DIR, exist_ok=True)
-
-
-async def create_enron_taskset() -> TaskInstanceSet:
-    ds_train = load_dataset(
-        "corbt/enron_emails_sample_questions",
-        split="train",
-        cache_dir=CACHE_DIR,
-    )
-    ds_test = load_dataset(
-        "corbt/enron_emails_sample_questions",
-        split="test",
-        cache_dir=CACHE_DIR,
-    )
-
-    def to_instance(row: dict, split: str) -> EnronTaskInstance:
-        impetus = Impetus(instructions=row["question"])
-        intent = Intent(
-            rubric={"goal": "Answer the question using the Enron emails."},
-            gold_trajectories=None,
-            gold_state_diff={"answer": row["answer"]},
-        )
-        metadata = EnronTaskInstanceMetadata(
-            split=split,
-            email_count=len(row["message_ids"]),
-            message_ids=row["message_ids"],
-        )
-        return EnronTaskInstance(
-            id=uuid4(),
-            impetus=impetus,
-            intent=intent,
-            metadata=metadata,
-            is_reproducible=True,
-            initial_engine_snapshot=row,
-        )
-
-    train_instances = [to_instance(r, "train") for r in ds_train]
-    test_instances = [to_instance(r, "test") for r in ds_test]
-
-    split_info = SplitInfo(
-        val_instance_ids=set(),
-        test_instance_ids={inst.id for inst in test_instances},
-        _is_split_defined=True,
-    )
-
-    return TaskInstanceSet(
-        name="Enron-QA",
-        description="QA over Enron email dataset sample.",
-        instances=train_instances + test_instances,
-        split_info=split_info,
-    )
-
-
-# quick sanity check ----------------------------------------------------------
-if __name__ == "__main__":
-
-    async def _main():
-        ts = await create_enron_taskset()
-        print(f"{len(ts.instances)} instances built.")
-
-    asyncio.run(_main())