open-data-sci 0.1.0__py3-none-any.whl

opendatasci/agents/turn_memory.py

@@ -0,0 +1,124 @@
+"""AgentLoopCompactor: LLM-based in-turn context compaction for ReAct loops."""
+
+import logging
+from typing import Any
+
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from langchain_core.messages.utils import count_tokens_approximately
+from pydantic import BaseModel
+
+from opendatasci._utils.langchain_utils import is_ongoing_turn, render_turn
+from opendatasci.prompts.prompt_templates import MIDTURN_COMPACTOR_SYSTEM_PROMPT
+
+logger = logging.getLogger(__name__)
+
+
+class _CompactedAgentLoop(BaseModel):
+    content: str
+
+    def to_message(self) -> SystemMessage:
+        """Return the compacted steps as a SystemMessage for injection into the turn."""
+        return SystemMessage(
+            content=f"<compacted_agent_loop>\n{self.content}\n</compacted_agent_loop>"
+        )
+
+
+class AgentLoopCompactor:
+    """Reduces an in-progress agent turn to fit within the token budget.
+
+    When a turn's intermediate steps exceed the configured threshold, the
+    compactor summarises them before the next LLM call.  Only the prompt
+    presented to the model is affected; the stored graph state is unchanged.
+    """
+
+    def __init__(self, llm: Any) -> None:
+        self._llm = llm
+        self._structured_llm = llm.with_structured_output(_CompactedAgentLoop)
+        self._system_prompt = MIDTURN_COMPACTOR_SYSTEM_PROMPT
+
+    def estimate_tokens(self, messages: list[BaseMessage]) -> int:
+        """Return an approximate token count for *messages*."""
+        return count_tokens_approximately(messages)
+
+    async def compact(self, turn: list[BaseMessage]) -> list[BaseMessage]:
+        """Compact *turn* by summarising intermediate steps and return the reduced message list.
+
+        The original user message and the most recent agent step are always preserved;
+        only the intermediate steps between them are replaced by a summary.
+
+        Returns *turn* unchanged if it is not an ongoing turn, if there are no
+        intermediate steps to compact, or if the LLM call fails.
+        """
+        if not is_ongoing_turn(turn):
+            return turn
+
+        # Find the last AIMessage in the turn.
+        last_ai_idx = -1
+        for i in range(len(turn) - 1, -1, -1):
+            if isinstance(turn[i], AIMessage):
+                last_ai_idx = i
+                break
+
+        if last_ai_idx == -1:
+            return turn
+
+        intermediate = turn[1:last_ai_idx]
+        if not intermediate:
+            return turn
+
+        try:
+            result: _CompactedAgentLoop = await self._structured_llm.ainvoke(
+                [
+                    SystemMessage(content=self._system_prompt),
+                    HumanMessage(content=render_turn(intermediate)),
+                ]
+            )
+
+            user_message: BaseMessage = turn[0]
+            compaction_message: BaseMessage = result.to_message()
+            rest_of_turn: list[BaseMessage] = turn[last_ai_idx:]
+            return [user_message, compaction_message, *rest_of_turn]
+        except Exception:
+            logger.exception("AgentLoopCompactor LLM call failed; return uncompacted turn")
+            return turn
+
+
+class TurnRewinder:
+    """Removes turns from a conversation history."""
+
+    def rewind_last_turn(
+        self,
+        chat_history: list[BaseMessage],
+        keep_user_message: bool = False,
+    ) -> list[BaseMessage]:
+        """Remove the last turn from *chat_history*.
+
+        The last turn is rewound regardless of whether it has completed — an
+        in-progress turn (started but no final agent response yet) is treated
+        the same as a completed one.
+
+        By default the human message that opened the turn is also dropped.
+        Pass ``keep_user_message=True`` to retain it and drop only the agent
+        response and any intermediate tool messages.
+
+        Args:
+            chat_history: The full conversation message list.
+            keep_user_message: When ``True``, the user message that opened
+                the last turn is retained. Defaults to ``False``.
+
+        Returns:
+            A new list with the last turn removed, or a copy of
+            *chat_history* unchanged if no turn start is found.
+        """
+        # Locate the HumanMessage that opened the last turn.
+        start_idx = -1
+        for i in range(len(chat_history) - 1, -1, -1):
+            if isinstance(chat_history[i], HumanMessage):
+                start_idx = i
+                break
+
+        if start_idx == -1:
+            return list(chat_history)
+
+        cut = start_idx + 1 if keep_user_message else start_idx
+        return list(chat_history[:cut])

opendatasci/configs.py

@@ -0,0 +1,275 @@
+"""Configuration for OpenDataSci."""
+
+from pathlib import Path
+from types import MappingProxyType
+from typing import List
+
+from pydantic import Field, model_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from opendatasci.models.providers import Provider
+from opendatasci.skills.local import _BUILTIN_SKILLS_DIRECTORY as _DEFAULT_BUILTIN_SKILLS_DIRECTORY
+
+DEFAULT_MODEL: MappingProxyType[Provider, str] = MappingProxyType(
+    {
+        Provider.ANTHROPIC: "claude-sonnet-4-6",
+        Provider.OPENAI: "gpt-5.5",
+        Provider.BEDROCK: "us.anthropic.claude-sonnet-4-6",
+        Provider.GEMINI: "gemini-2.5-pro",
+        Provider.VERTEXAI: "gemini-2.5-pro",
+        Provider.AZURE: "gpt-4o",
+        Provider.OLLAMA: "llama3.2:3b",
+        Provider.OPENAI_COMPATIBLE_SERVER: "meta-llama/Llama-3.2-3B-Instruct",
+    }
+)
+
+DEFAULT_SECONDARY_MODEL: MappingProxyType[Provider, str] = MappingProxyType(
+    {
+        Provider.ANTHROPIC: "claude-haiku-4-5",
+        Provider.OPENAI: "gpt-5.4-mini",
+        Provider.BEDROCK: "us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        Provider.GEMINI: "gemini-2.5-flash",
+        Provider.VERTEXAI: "gemini-2.5-flash",
+        Provider.AZURE: "gpt-4o-mini",
+        Provider.OLLAMA: "llama3.2:3b",
+        Provider.OPENAI_COMPATIBLE_SERVER: "meta-llama/Llama-3.2-3B-Instruct",
+    }
+)
+
+
+class OpenDataSciConfig(BaseSettings):
+    """Configuration for OpenDataSci.
+
+    All fields can be set via environment variables (names shown in parentheses)
+    or a ``.env`` file.  Pass an instance to :func:`create_agent` or
+    :class:`Agent` to apply custom settings.
+
+    Attributes:
+        provider:        LLM provider for the primary model.  One of
+                         ``"anthropic"``, ``"openai"``, ``"bedrock"``,
+                         ``"gemini"``, ``"vertexai"``, ``"azure"``,
+                         ``"ollama"``, ``"openai_compatible_server"`` (any
+                         self-hosted OpenAI-compatible server, e.g. vLLM).
+        model:           Provider-specific model identifier.  Falls back to a
+                         sensible default per provider when not set.
+        secondary_provider: LLM provider for the secondary (auxiliary) model.
+                         Defaults to ``provider`` when not set.  Set to a
+                         different provider to mix backends — e.g. Anthropic
+                         for the primary model and OpenAI for summarisation.
+        secondary_model: Model identifier for lightweight tasks such as memory
+                         summarisation.  Falls back to a sensible default per
+                         provider when not set.
+        anthropic_api_key: API key for Anthropic (``ANTHROPIC_API_KEY``).
+        openai_api_key:    API key for OpenAI and OpenAI-compatible servers
+                           (``OPENAI_API_KEY``).
+        google_api_key:    API key for Google Gemini (``GOOGLE_API_KEY``).
+        azure_api_key:     API key for Azure OpenAI (``AZURE_OPENAI_API_KEY``).
+                           Omit when using service-principal auth instead.
+        aws_region:        AWS region for Bedrock (``REGION``).
+        google_cloud_project:  GCP project ID for Vertex AI
+                           (``GOOGLE_CLOUD_PROJECT``).
+        google_cloud_location: Vertex AI region / location
+                           (``GOOGLE_CLOUD_LOCATION``).
+        azure_endpoint:    Azure OpenAI resource endpoint URL
+                           (``AZURE_OPENAI_ENDPOINT``).  Required when
+                           ``provider`` is ``"azure"``.
+        azure_api_version: Azure OpenAI API version
+                           (``AZURE_OPENAI_API_VERSION``).  Defaults to
+                           ``"2025-01-01-preview"``.
+        llm_server_base_url: Base URL for self-hosted providers
+                           (``LLM_SERVER_BASE_URL``).  Required for
+                           ``"ollama"`` and ``"openai_compatible_server"``;
+                           falls back to ``http://localhost:11434`` and
+                           ``http://localhost:8000/v1`` respectively when
+                           not set.
+        temperature:     LLM sampling temperature.  Ignored for Anthropic and
+                         Bedrock when extended thinking is active (those
+                         providers require temperature ``1``).
+        thinking_budget: Token budget for extended thinking / reasoning
+                         (Anthropic and Bedrock only).
+        name:            Display name of the agent.  Defaults to ``"Sai"``.
+        mcp_servers: List of MCP server URLs the agent may connect to
+                         (``MCP_SERVERS``).
+        skills_directory: Path to a directory of custom skill files
+                         (``SKILLS_DIRECTORY``).  Loaded in addition to the
+                         built-in skills; custom skills override built-ins of
+                         the same name.
+        builtin_skills_directory: Path to the built-in skills bundled with
+                         the package (``BUILTIN_SKILLS_DIRECTORY``).  Override
+                         only if you need to replace the defaults entirely.
+        extra_web_domains: Additional hostnames the ``fetch_url`` tool may
+                         retrieve, on top of the built-in allowlist.
+                         Example: ``["internal.corp"]``.
+        override_web_domains: When set, *replaces* the built-in domain
+                         allowlist entirely.  ``extra_web_domains`` is still
+                         applied on top.  Use ``[]`` to block all domains.
+        worker_timeout_seconds: Maximum seconds to wait for all spawned
+                         workers to finish.  ``None`` disables the timeout.
+                         Defaults to ``300.0`` (5 minutes).
+        midturn_compaction_threshold: Token count after which the agent's
+                         context is compacted mid-turn (during a single turn's
+                         reasoning/acting loop).  Must be strictly positive.
+                         Defaults to ``80000``.
+        local_code_exec_timeout: Maximum seconds allowed for a single
+                         code-execution run in the local sandbox
+                         (``CODE_EXEC_TIMEOUT``).  Defaults to
+                         ``1800`` (30 minutes).
+    Cloud authentication (environment variables read by the underlying SDKs):
+
+    **AWS Bedrock** — boto3 credential chain (pick one):
+
+    - Long-lived IAM key: ``AWS_ACCESS_KEY_ID`` + ``AWS_SECRET_ACCESS_KEY``
+    - Temporary STS credentials: add ``AWS_SESSION_TOKEN`` to the above.
+    - EC2 / ECS / Lambda: credentials are fetched automatically from the
+      instance metadata service; no env vars required.
+
+    **Google Vertex AI** — Application Default Credentials chain (pick one):
+
+    - Service account JSON key: set ``GOOGLE_APPLICATION_CREDENTIALS`` to
+      the path of the key file; also set ``GOOGLE_CLOUD_PROJECT``.
+    - User credentials: run ``gcloud auth application-default login``.
+    - Cloud Run / GCE / GKE: credentials are fetched automatically;
+      ``GOOGLE_CLOUD_PROJECT`` is still required.
+
+    **Azure OpenAI** — API key *or* service principal (not both):
+
+    - API key: ``AZURE_OPENAI_API_KEY``
+    - Service principal (requires ``pip install 'open-data-sci[azure]'``):
+      set ``AZURE_TENANT_ID``, ``AZURE_CLIENT_ID``, and
+      ``AZURE_CLIENT_SECRET``.
+    """
+
+    model_config = SettingsConfigDict(
+        frozen=True,
+        populate_by_name=True,
+        env_ignore_empty=True,
+        env_file=".env",
+    )
+
+    # ── Model selection ───────────────────────────────────────────────────────
+    provider: Provider = Field(default=Provider.ANTHROPIC, alias="PROVIDER")
+    model: str = Field(default=DEFAULT_MODEL[Provider.ANTHROPIC], alias="MODEL")
+    secondary_provider: Provider = Field(default=Provider.ANTHROPIC, alias="SECONDARY_PROVIDER")
+    secondary_model: str = Field(
+        default=DEFAULT_SECONDARY_MODEL[Provider.ANTHROPIC], alias="SECONDARY_MODEL"
+    )
+
+    # ── Per-provider API keys (loaded from env via alias) ─────────────────────
+    anthropic_api_key: str | None = Field(default=None, alias="ANTHROPIC_API_KEY")
+    openai_api_key: str | None = Field(default=None, alias="OPENAI_API_KEY")
+    google_api_key: str | None = Field(default=None, alias="GOOGLE_API_KEY")
+    azure_api_key: str | None = Field(default=None, alias="AZURE_OPENAI_API_KEY")
+
+    # ── Cloud region / location ───────────────────────────────────────────────
+    aws_region: str | None = Field(default=None, alias="REGION")
+    google_cloud_project: str | None = Field(default=None, alias="GOOGLE_CLOUD_PROJECT")
+    google_cloud_location: str | None = Field(default=None, alias="GOOGLE_CLOUD_LOCATION")
+
+    # ── Azure-specific ────────────────────────────────────────────────────────
+    azure_endpoint: str | None = Field(default=None, alias="AZURE_OPENAI_ENDPOINT")
+    azure_api_version: str = Field(default="2025-01-01-preview", alias="AZURE_OPENAI_API_VERSION")
+
+    # ── Self-hosted endpoint (Ollama / OpenAI-compatible server) ──────────────
+    llm_server_base_url: str | None = Field(default=None, alias="LLM_SERVER_BASE_URL")
+
+    # ── Sampling & reasoning ──────────────────────────────────────────────────
+    temperature: float = Field(default=0.0, alias="TEMPERATURE")
+    thinking_budget: int = Field(default=8192, alias="THINKING_BUDGET")
+
+    # ── Agent Customization ───────────────────────────────────────────────────────
+    name: str = Field(default="Sai", alias="NAME")
+
+    # ── MCP ───────────────────────────────────────────────────────────
+    mcp_servers: List[str] = Field(default_factory=list, alias="MCP_SERVERS")
+
+    # ── Web access ───────────────────────────────────────────────────────────────
+    extra_web_domains: List[str] = Field(default_factory=list, alias="EXTRA_FETCH_DOMAINS")
+    override_web_domains: List[str] | None = None
+
+    # ── Context management ───────────────────────────────────────────────────────
+    midturn_compaction_threshold: int = Field(
+        default=64000,
+        alias="MIDTURN_COMPACTION_THRESHOLD",
+        gt=0,
+        description="Token count after which the agent's context is compacted mid-turn. Only applies in execution mode.",
+    )
+
+    # ── Skills ───────────────────────────────────────────────────────────────────
+    skills_directory: Path | None = Field(
+        default=None,
+        alias="SKILLS_DIRECTORY",
+    )
+    builtin_skills_directory: Path = Field(
+        default=_DEFAULT_BUILTIN_SKILLS_DIRECTORY,
+        alias="BUILTIN_SKILLS_DIRECTORY",
+    )
+
+    # ── Worker configuration ───────────────────────────────────────────────────────
+    worker_timeout_seconds: float | None = Field(
+        default=300.0,
+        alias="WORKER_TIMEOUT_SECONDS",
+    )
+
+    # ── Sandbox ───────────────────────────────────────────────────────────────────
+    local_code_exec_timeout: int = Field(
+        default=1800,  # 30 minutes
+        alias="CODE_EXEC_TIMEOUT",
+        gt=0,
+        description="Maximum seconds allowed for a single local code-execution run.",
+    )
+
+    @model_validator(mode="after")
+    def _validate_providers(self) -> "OpenDataSciConfig":
+        if self.provider not in DEFAULT_MODEL:
+            supported = ", ".join(f"'{p}'" for p in sorted(DEFAULT_MODEL))
+            raise ValueError(
+                f"Unknown provider '{self.provider}'. Supported providers: {supported}."
+            )
+        if self.secondary_provider is not None and self.secondary_provider not in DEFAULT_MODEL:
+            supported = ", ".join(f"'{p}'" for p in sorted(DEFAULT_MODEL))
+            raise ValueError(
+                f"Unknown secondary model provider '{self.secondary_provider}'. "
+                f"Supported providers: {supported}."
+            )
+        return self
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> "OpenDataSciConfig":
+        """Load an ``OpenDataSciConfig`` from a YAML file.
+
+        The file must be a YAML mapping whose keys match ``OpenDataSciConfig``
+        field names.  Unknown keys raise ``ValueError`` with a descriptive
+        message.  Environment variables are still applied for any field not
+        present in the file.
+
+        Raises:
+            ImportError: If PyYAML is not installed.
+            ValueError: If the file does not contain a mapping, or contains
+                unknown field names.
+        """
+        try:
+            import yaml  # type: ignore[import-untyped]
+        except ImportError as exc:
+            raise ImportError(
+                "PyYAML is required to load YAML config files. "
+                "Install it with: pip install pyyaml"
+            ) from exc
+
+        with open(path) as fh:
+            data = yaml.safe_load(fh) or {}
+
+        if not isinstance(data, dict):
+            raise ValueError(
+                f"YAML config at '{path}' must be a mapping (key: value pairs), "
+                f"got {type(data).__name__}"
+            )
+
+        valid_fields = set(cls.model_fields.keys())
+        unknown = set(data) - valid_fields
+        if unknown:
+            raise ValueError(
+                f"Unknown fields in YAML config '{path}': {', '.join(sorted(unknown))}. "
+                f"Valid fields: {', '.join(sorted(valid_fields))}"
+            )
+
+        return cls(**data)

opendatasci/context/__init__.py

@@ -0,0 +1,7 @@
+from opendatasci.context.base import BaseContextStore
+from opendatasci.context.local import LocalContextStore
+
+__all__ = [
+    "BaseContextStore",
+    "LocalContextStore",
+]

opendatasci/context/base.py

@@ -0,0 +1,56 @@
+"""Abstract base classes for the agent's context stores."""
+
+from abc import ABC, abstractmethod
+from contextlib import AbstractAsyncContextManager
+from pathlib import Path
+from typing import Self
+
+
+class BaseContextStore(ABC):
+    """Context store for dataset notes, profile cards, and session plans.
+
+    Dataset notes and profile cards persist across agent sessions and are keyed
+    by dataset path.  Plans are scoped to a session and keyed by ``session_id``.
+    """
+
+    @abstractmethod
+    def session(self) -> AbstractAsyncContextManager[Self]:
+        """Return an async context manager that manages this store's lifecycle.
+
+        Implementations should yield ``self``, performing any required setup on
+        entry and teardown (flushing, releasing resources, etc.) on exit.
+        """
+
+    @property
+    @abstractmethod
+    def root(self) -> Path:
+        """Return the root path of the context store (e.g. the ``.opendatasci`` directory)."""
+
+    @abstractmethod
+    async def read_dataset_info(self, dataset_path: str) -> str:
+        """Return combined dataset info: profile card (if any) + session notes."""
+
+    @abstractmethod
+    async def update_dataset_info(
+        self,
+        dataset_path: str,
+        update: str,
+        merge: bool = True,
+    ) -> str:
+        """Persist dataset notes and return the path to the stored notes file."""
+
+    @abstractmethod
+    async def get_profile_info(self, dataset_path: str) -> tuple[str, str, str | None]:
+        """Return ``(resolved_path_str, hash_hex, existing_profile_or_None)``."""
+
+    @abstractmethod
+    def save_dataset_profile(self, hash_hex: str, content: str) -> None:
+        """Persist a completed profile card for *hash_hex*."""
+
+    @abstractmethod
+    def current_plan(self, session_id: str) -> str | None:
+        """Return the most recent plan for *session_id*, or ``None``."""
+
+    @abstractmethod
+    def save_plan(self, session_id: str, plan: str) -> None:
+        """Persist *plan* for *session_id*."""