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

opendatasci/context/local.py

@@ -0,0 +1,236 @@
+"""File-based context stores backed by the workspace's ``.opendatasci`` directory."""
+
+import datetime
+import logging
+from contextlib import asynccontextmanager
+from datetime import datetime as _datetime
+from datetime import timezone
+from pathlib import Path
+from typing import AsyncGenerator, Self
+
+from opendatasci._utils.hash_utils import hash_path
+from opendatasci.context.base import BaseContextStore
+
+logger = logging.getLogger(__name__)
+
+OPENDATASCI_DIRNAME = ".opendatasci"
+
+_NOTES_DIR = "dataset_notes"
+_PROFILES_DIR = "dataset_profiling"
+_PLANS_DIR = "plans"
+
+
+class LocalContextStore(BaseContextStore):
+    """File-based context store for a single local workspace.
+
+    Persists dataset notes and profile cards (keyed by dataset path) as well as
+    session plans (keyed by ``session_id``).  All data lives under the
+    workspace's ``.opendatasci`` directory.
+
+    Args:
+        workspace_path: Root directory of the active workspace.  Relative
+            dataset paths are resolved against this directory.
+    """
+
+    def __init__(self, workspace_path: Path) -> None:
+        self._workspace_path = workspace_path
+        self._root = workspace_path / OPENDATASCI_DIRNAME
+        self._current_plan: str | None = None
+
+    # ── BaseContextStore: session ────────────────────────────────────
+
+    @asynccontextmanager
+    async def session(self) -> AsyncGenerator[Self, None]:
+        self._root.mkdir(parents=True, exist_ok=True)
+        yield self
+
+    # ── BaseContextStore: root ──────────────────────────────────────
+
+    @property
+    def root(self) -> Path:
+        return self._root
+
+    # ── Internal storage paths ───────────────────────────────────────────────
+
+    @property
+    def _notes_root(self) -> Path:
+        return self.root / _NOTES_DIR
+
+    @property
+    def _profiles_root(self) -> Path:
+        return self.root / _PROFILES_DIR
+
+    @property
+    def _plans_root(self) -> Path:
+        return self.root / _PLANS_DIR
+
+    # ── Internal notes storage ────────────────────────────────────────────────
+
+    def _find(self, hash_hex: str) -> Path | None:
+        """Locate an existing notes file for *hash_hex* under any date prefix."""
+        if self._notes_root.exists():
+            for p in self._notes_root.rglob(f"{hash_hex}.md"):
+                return p
+        return None
+
+    def _resolve_notes_path(self, hash_hex: str) -> Path:
+        """Reuse an existing file for *hash_hex*, or build a new date-keyed path."""
+        existing = self._find(hash_hex)
+        if existing is not None:
+            return existing
+        today = datetime.date.today()
+        return (
+            self._notes_root
+            / str(today.year)
+            / f"{today.month:02d}"
+            / f"{today.day:02d}"
+            / f"{hash_hex}.md"
+        )
+
+    def _load_notes(self, hash_hex: str) -> str | None:
+        p = self._find(hash_hex)
+        return p.read_text(encoding="utf-8") if p is not None else None
+
+    def _save_notes(self, hash_hex: str, content: str) -> None:
+        path = self._resolve_notes_path(hash_hex)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(content, encoding="utf-8")
+
+    def _notes_file_path(self, hash_hex: str) -> str:
+        return str(self._resolve_notes_path(hash_hex))
+
+    # ── Internal profile storage ──────────────────────────────────────────────
+
+    def _load_profile(self, hash_hex: str) -> str | None:
+        p = self._profiles_root / f"{hash_hex}.md"
+        return p.read_text(encoding="utf-8") if p.exists() else None
+
+    # ── BaseContextStore interface ──────────────────────────────────
+
+    async def read_dataset_info(self, dataset_path: str) -> str:
+        """Return combined dataset info: profile card (if any) + session notes.
+
+        The returned Markdown string has clearly labelled sections:
+
+        ``# DATASET PROFILING`` — auto-generated stats (shape, dtypes, etc.)
+        ``# DATASET NOTES``     — persistent notes written by the agent
+
+        If no profile exists, only the notes section is included.
+        If no notes exist yet, a placeholder scaffold is returned.
+        """
+        path = self._resolve_dataset_path(dataset_path)
+        hash_hex = await hash_path(path)
+
+        notes = self._load_notes(hash_hex)
+        if notes is None:
+            notes = f"# DATASET NOTES: {path.name}\n\n_No notes recorded yet._\n"
+        else:
+            notes = "# DATASET NOTES\n\n" + notes
+
+        profile = self._load_profile(hash_hex)
+        if profile is not None:
+            return "# DATASET PROFILING\n\n" + profile.rstrip() + "\n\n---\n\n" + notes
+        return notes
+
+    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.
+
+        Only notes are modified — the profile card (if any) is never touched.
+        """
+        path = self._resolve_dataset_path(dataset_path)
+        hash_hex = await hash_path(path)
+
+        if merge:
+            existing = self._load_notes(hash_hex)
+            new_content = (
+                existing.rstrip() + "\n\n" + update.strip() + "\n"
+                if existing
+                else update.strip() + "\n"
+            )
+        else:
+            new_content = update.strip() + "\n"
+
+        self._save_notes(hash_hex, new_content)
+        return self._notes_file_path(hash_hex)
+
+    async def get_profile_info(self, dataset_path: str) -> tuple[str, str, str | None]:
+        """Return ``(resolved_path_str, hash_hex, existing_profile_or_None)``.
+
+        Used by the ``profile_dataset`` tool to check for a cached card before
+        running the profiling sandbox pass.
+        """
+        path = self._resolve_dataset_path(dataset_path)
+        hash_hex = await hash_path(path)
+        return str(path), hash_hex, self._load_profile(hash_hex)
+
+    def save_dataset_profile(self, hash_hex: str, content: str) -> None:
+        """Persist a completed profile card for *hash_hex*."""
+        self._profiles_root.mkdir(parents=True, exist_ok=True)
+        (self._profiles_root / f"{hash_hex}.md").write_text(content, encoding="utf-8")
+
+    # ── Private helpers ───────────────────────────────────────────────────────
+
+    def _resolve_dataset_path(self, dataset_path: str) -> Path:
+        p = Path(dataset_path)
+        path = (self._workspace_path / p).resolve() if not p.is_absolute() else p.resolve()
+        if not path.exists():
+            raise FileNotFoundError(f"Dataset path does not exist: {path}")
+        return path
+
+    # ── BaseContextStore: plans ───────────────────────────────────────────────
+
+    def current_plan(self, session_id: str) -> str | None:
+        """Return the most recent plan for this session.
+
+        Resolves the plan by reading the latest ``{session_id}_*.txt`` file from
+        disk so the plan survives process restarts.  Falls back to the in-memory
+        cache when no file exists yet or a disk read fails.
+        """
+        if self._plans_root.exists():
+            files = sorted(self._plans_root.glob(f"{session_id}_*.txt"))
+            if files:
+                try:
+                    return files[-1].read_text(encoding="utf-8")
+                except OSError:
+                    logger.warning("Could not read plan file: %s", files[-1], exc_info=True)
+        return self._current_plan
+
+    def save_plan(self, session_id: str, plan: str) -> None:
+        """Persist *plan* to disk and prune stale files.
+
+        Also updates the in-memory cache used as a fallback when no plan file
+        can be read back.
+        """
+        self._current_plan = plan
+        self._plans_root.mkdir(parents=True, exist_ok=True)
+        stamp = _datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+        path = self._plans_root / f"{session_id}_{stamp}.txt"
+        try:
+            path.write_text(plan, encoding="utf-8")
+        except OSError:
+            logger.warning("Could not write plan file: %s", path, exc_info=True)
+            return
+        self.prune()
+
+    def prune(self) -> None:
+        """Keep only the most recent plan file per session_id."""
+        if not self._plans_root.exists():
+            return
+        by_session: dict[str, list[Path]] = {}
+        for p in self._plans_root.glob("*.txt"):
+            sid = p.stem.partition("_")[0]
+            if not sid:
+                continue
+            by_session.setdefault(sid, []).append(p)
+        for files in by_session.values():
+            files.sort(key=lambda f: f.name)
+            for stale in files[:-1]:
+                try:
+                    stale.unlink()
+                except OSError:
+                    logger.exception("Could not delete stale plan file: %s", stale)

opendatasci/models/__init__.py

@@ -0,0 +1,7 @@
+from opendatasci.models.factory import (
+    create_model,
+    create_secondary_model,
+)
+from opendatasci.models.providers import Provider
+
+__all__ = ["create_model", "create_secondary_model", "Provider"]

opendatasci/models/anthropic.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from langchain_core.language_models import BaseChatModel
+
+from opendatasci.configs import OpenDataSciConfig
+
+
+def create_anthropic_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a ``ChatAnthropic`` model with extended thinking enabled."""
+    try:
+        from langchain_anthropic import ChatAnthropic
+    except ImportError as exc:
+        raise ValueError("langchain-anthropic is not installed.") from exc
+    return ChatAnthropic(
+        model=config.model,
+        api_key=config.anthropic_api_key,
+        # Temperature must be 1 when extended thinking is enabled.
+        temperature=1,
+        max_tokens=16000,
+        thinking={"type": "enabled", "budget_tokens": config.thinking_budget},
+    )
+
+
+def create_anthropic_secondary_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a cheap ``ChatAnthropic`` model for auxiliary tasks (thinking disabled)."""
+    try:
+        from langchain_anthropic import ChatAnthropic
+    except ImportError as exc:
+        raise ValueError("langchain-anthropic is not installed.") from exc
+    return ChatAnthropic(
+        model=config.secondary_model,
+        api_key=config.anthropic_api_key,
+        temperature=0,
+        max_tokens=1000,
+    )
+
+
+def cached_system_prompt(prompt: str) -> list[dict[str, Any]]:
+    """Wrap *prompt* with Anthropic's ephemeral cache breakpoint."""
+    return [{"type": "text", "text": prompt, "cache_control": {"type": "ephemeral"}}]

opendatasci/models/aws.py

@@ -0,0 +1,86 @@
+from typing import Any
+
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessageChunk
+from langchain_core.outputs import ChatGenerationChunk
+
+from opendatasci.configs import OpenDataSciConfig
+
+try:
+    from langchain_aws import ChatBedrockConverse as _BedrockBase
+except ImportError:
+    _BedrockBase = None
+
+
+def _strip_list_usage_fields(chunk: ChatGenerationChunk) -> ChatGenerationChunk:
+    """Remove list-valued fields from a chunk's usage_metadata.
+
+    Bedrock's TokenUsage includes a ``cacheDetails`` field (type: list) which
+    langchain_core's ``_dict_int_op`` cannot combine — it only accepts int/dict
+    values.  Stripping the field here prevents the ValueError that would
+    otherwise surface as a spurious error event at the end of each streamed
+    response when prompt caching is active.
+    """
+    msg = chunk.message
+    if not (isinstance(msg, AIMessageChunk) and msg.usage_metadata):
+        return chunk
+    if not any(isinstance(v, list) for v in msg.usage_metadata.values()):
+        return chunk
+    cleaned = {k: v for k, v in msg.usage_metadata.items() if not isinstance(v, list)}
+    return ChatGenerationChunk(
+        message=msg.model_copy(update={"usage_metadata": cleaned}),
+        generation_info=chunk.generation_info,
+    )
+
+
+if _BedrockBase is not None:
+
+    class _CustomBedrockConverse(_BedrockBase):  # type: ignore[misc]
+        def _stream(
+            self, messages: Any, stop: Any = None, run_manager: Any = None, **kwargs: Any
+        ) -> Any:
+            for chunk in super()._stream(messages, stop, run_manager, **kwargs):
+                yield _strip_list_usage_fields(chunk)
+else:
+    _CustomBedrockConverse = None  # type: ignore[assignment,misc]
+
+
+def create_bedrock_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a Bedrock Converse model with adaptive thinking and streaming enabled."""
+    if _CustomBedrockConverse is None:
+        raise ValueError("langchain-aws is not installed.")
+    return _CustomBedrockConverse(
+        model=config.model,
+        region_name=config.aws_region,
+        # Temperature must be 1 when extended thinking is enabled.
+        temperature=1,
+        max_tokens=16000,
+        additional_model_request_fields={
+            "thinking": {"type": "adaptive"},
+            "output_config": {"effort": "medium"},
+        },
+        # langchain-aws's set_disable_streaming validator only recognises
+        # "claude-3" model IDs as streaming-capable.  Claude 4 models
+        # (claude-sonnet-4-6, etc.) match none of the patterns and default to
+        # disable_streaming=True, causing every response to arrive as a single
+        # non-streamed chunk.  Claude 4 supports Bedrock ConverseStream with
+        # tools, so we override the auto-detection explicitly.
+        disable_streaming=False,
+    )
+
+
+def create_bedrock_secondary_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a cheap Bedrock model for auxiliary tasks (thinking disabled)."""
+    if _BedrockBase is None:
+        raise ValueError("langchain-aws is not installed.")
+    return _BedrockBase(  # type: ignore[no-any-return]
+        model=config.secondary_model,
+        region_name=config.aws_region,
+        temperature=0,
+        max_tokens=1000,
+    )
+
+
+def cached_system_prompt(prompt: str) -> list[dict[str, Any]]:
+    """Wrap *prompt* with a Bedrock Converse cache point breakpoint."""
+    return [{"type": "text", "text": prompt}, {"cachePoint": {"type": "default"}}]

opendatasci/models/factory.py

@@ -0,0 +1,179 @@
+"""
+LLM provider factory for the Agent.
+"""
+
+import asyncio
+import logging
+import random
+from typing import Any
+
+from langchain_core.language_models import BaseChatModel
+
+from opendatasci.configs import OpenDataSciConfig
+from opendatasci.models.providers import Provider
+
+_LOG = logging.getLogger(__name__)
+
+_TRANSIENT_KEYWORDS = frozenset(
+    [
+        "rate limit",
+        "ratelimit",
+        "rate_limit",
+        "429",
+        "too many requests",
+        "overloaded",
+        "overload",
+        "503",
+        "service unavailable",
+        "throttl",
+        "connection error",
+        "connection refused",
+        "connecterror",
+        "apiconnectionerror",
+    ]
+)
+
+_MAX_RETRY_WAIT = 60.0
+_DEFAULT_MAX_ATTEMPTS = 5
+
+
+def _is_transient(exc: Exception) -> bool:
+    text = (type(exc).__name__ + " " + str(exc)).lower()
+    return any(kw in text for kw in _TRANSIENT_KEYWORDS)
+
+
+class _RetryRunnable:
+    """Wraps a LangChain Runnable, retrying transient errors with exponential backoff."""
+
+    def __init__(self, runnable: Any, max_attempts: int = _DEFAULT_MAX_ATTEMPTS) -> None:
+        self._runnable = runnable
+        self._max_attempts = max_attempts
+
+    async def ainvoke(self, *args: Any, **kwargs: Any) -> Any:
+        for attempt in range(self._max_attempts):
+            try:
+                return await self._runnable.ainvoke(*args, **kwargs)
+            except Exception as exc:
+                if _is_transient(exc) and attempt < self._max_attempts - 1:
+                    wait = min(_MAX_RETRY_WAIT, (2**attempt) + random.random())
+                    _LOG.warning(
+                        "LLM transient error (attempt %d/%d), retrying in %.1fs: %s",
+                        attempt + 1,
+                        self._max_attempts,
+                        wait,
+                        exc,
+                    )
+                    await asyncio.sleep(wait)
+                else:
+                    raise
+
+
+def with_retry(runnable: Any, max_attempts: int = _DEFAULT_MAX_ATTEMPTS) -> _RetryRunnable:
+    """Wrap *runnable* so that transient LLM errors are retried with exponential backoff."""
+    return _RetryRunnable(runnable, max_attempts=max_attempts)
+
+
+def _supported_providers() -> str:
+    return ", ".join(f"'{p}'" for p in Provider)
+
+
+def create_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate the primary LLM for the agent and workers.
+
+    Thinking is enabled by default for providers that support it (Anthropic, Bedrock).
+    Falls back to the provider's default primary model when none is specified.
+    """
+    match config.provider:
+        case Provider.ANTHROPIC:
+            from opendatasci.models.anthropic import create_anthropic_model as _create
+
+            return _create(config)
+        case Provider.BEDROCK:
+            from opendatasci.models.aws import create_bedrock_model as _create
+
+            return _create(config)
+        case Provider.OPENAI:
+            from opendatasci.models.openai import create_openai_model as _create
+
+            return _create(config)
+        case Provider.GEMINI:
+            from opendatasci.models.google import create_gemini_model as _create
+
+            return _create(config)
+        case Provider.VERTEXAI:
+            from opendatasci.models.google import create_vertexai_model as _create
+
+            return _create(config)
+        case Provider.AZURE:
+            from opendatasci.models.microsoft import create_azure_model as _create
+
+            return _create(config)
+        case Provider.OLLAMA:
+            from opendatasci.models.local import create_ollama_model as _create
+
+            return _create(config)
+        case Provider.OPENAI_COMPATIBLE_SERVER:
+            from opendatasci.models.local import create_openai_compatible_model as _create
+
+            return _create(config)
+
+    raise ValueError(
+        f"Unknown provider '{config.provider}'. Supported providers: {_supported_providers()}."
+    )
+
+
+def create_secondary_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate the secondary (auxiliary) LLM for lightweight tasks such as summarization.
+
+    Thinking is disabled. Falls back to the provider's default secondary model when none is
+    specified.
+
+    When ``config.secondary_provider`` differs from ``config.provider``, a
+    lightweight config copy is used so the secondary model's factory resolves its
+    API key from its own environment variable rather than the primary provider's.
+    """
+    secondary_provider = config.secondary_provider
+
+    # When the secondary model uses a different provider, switch the provider
+    # field so the factory dispatches correctly.  Each provider reads its own
+    # dedicated API-key field, so no credential clearing is needed.
+    if secondary_provider != config.provider:
+        config = config.model_copy(update={"provider": secondary_provider})
+
+    match secondary_provider:
+        case Provider.ANTHROPIC:
+            from opendatasci.models.anthropic import create_anthropic_secondary_model as _create
+
+            return _create(config)
+        case Provider.BEDROCK:
+            from opendatasci.models.aws import create_bedrock_secondary_model as _create
+
+            return _create(config)
+        case Provider.OPENAI:
+            from opendatasci.models.openai import create_openai_secondary_model as _create
+
+            return _create(config)
+        case Provider.GEMINI:
+            from opendatasci.models.google import create_gemini_secondary_model as _create
+
+            return _create(config)
+        case Provider.VERTEXAI:
+            from opendatasci.models.google import create_vertexai_secondary_model as _create
+
+            return _create(config)
+        case Provider.AZURE:
+            from opendatasci.models.microsoft import create_azure_secondary_model as _create
+
+            return _create(config)
+        case Provider.OLLAMA:
+            from opendatasci.models.local import create_ollama_secondary_model as _create
+
+            return _create(config)
+        case Provider.OPENAI_COMPATIBLE_SERVER:
+            from opendatasci.models.local import create_openai_compatible_secondary_model as _create
+
+            return _create(config)
+
+    raise ValueError(
+        f"Unknown provider '{secondary_provider}'. Supported providers: {_supported_providers()}."
+    )

opendatasci/models/google.py

@@ -0,0 +1,79 @@
+from langchain_core.language_models import BaseChatModel
+
+from opendatasci.configs import OpenDataSciConfig
+
+
+def create_gemini_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a ``ChatGoogleGenerativeAI`` model via the Gemini API."""
+    try:
+        from langchain_google_genai import ChatGoogleGenerativeAI
+    except ImportError as exc:
+        raise ValueError(
+            "langchain-google-genai is not installed. Run: pip install 'open-data-sci[gemini]'"
+        ) from exc
+    return ChatGoogleGenerativeAI(  # type: ignore[no-any-return]
+        model=config.model,
+        google_api_key=config.google_api_key,
+        temperature=config.temperature,
+    )
+
+
+def create_gemini_secondary_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a cheap Gemini model for auxiliary tasks."""
+    try:
+        from langchain_google_genai import ChatGoogleGenerativeAI
+    except ImportError as exc:
+        raise ValueError(
+            "langchain-google-genai is not installed. Run: pip install 'open-data-sci[gemini]'"
+        ) from exc
+    return ChatGoogleGenerativeAI(  # type: ignore[no-any-return]
+        model=config.secondary_model,
+        google_api_key=config.google_api_key,
+        temperature=0,
+        max_output_tokens=1000,
+    )
+
+
+def create_vertexai_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a ``ChatVertexAI`` model via Google Cloud Vertex AI."""
+    try:
+        from langchain_google_vertexai import ChatVertexAI  # type: ignore[import-not-found]
+    except ImportError as exc:
+        raise ValueError(
+            "langchain-google-vertexai is not installed. Run: pip install 'open-data-sci[gcp]'"
+        ) from exc
+    return ChatVertexAI(  # type: ignore[no-any-return]
+        model=config.model,
+        project=config.google_cloud_project,
+        location=config.google_cloud_location,
+        temperature=config.temperature,
+    )
+
+
+def create_vertexai_secondary_model(config: OpenDataSciConfig) -> BaseChatModel:
+    """Instantiate a cheap Vertex AI model for auxiliary tasks."""
+    try:
+        from langchain_google_vertexai import ChatVertexAI
+    except ImportError as exc:
+        raise ValueError(
+            "langchain-google-vertexai is not installed. Run: pip install 'open-data-sci[gcp]'"
+        ) from exc
+    return ChatVertexAI(  # type: ignore[no-any-return]
+        model=config.secondary_model,
+        project=config.google_cloud_project,
+        location=config.google_cloud_location,
+        temperature=0,
+        max_output_tokens=1000,
+    )
+
+
+def cached_system_prompt(prompt: str) -> str:
+    """Return the system prompt unchanged.
+
+    Gemini 2.5+ models perform implicit context caching automatically for
+    prompts above the per-model minimum (>= 1024 tokens for Flash, >= 4096 for
+    Pro), keying off the request's leading prefix. No client-side cache
+    markers are required, and explicit `cached_content` setup is intentionally
+    not used here to keep the model factory side-effect free.
+    """
+    return prompt