aixtools 0.1.9__tar.gz → 0.1.11__tar.gz

{aixtools-0.1.9 → aixtools-0.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.1.9
+Version: 0.1.11
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown
@@ -16,6 +16,7 @@ Requires-Dist: langchain-chroma>=0.2.3
 Requires-Dist: langchain-ollama>=0.3.2
 Requires-Dist: langchain-openai>=0.3.14
 Requires-Dist: mcp>=1.11.0
+Requires-Dist: mypy>=1.18.2
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: pydantic-ai>=0.4.10
 Requires-Dist: pylint>=3.3.7

{aixtools-0.1.9 → aixtools-0.1.11}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.9'
-__version_tuple__ = version_tuple = (0, 1, 9)
+__version__ = version = '0.1.11'
+__version_tuple__ = version_tuple = (0, 1, 11)
 
-__commit_id__ = commit_id = 'g3e272f9d5'
+__commit_id__ = commit_id = 'g32157e2f8'

aixtools-0.1.11/aixtools/compliance/__init__.py

@@ -0,0 +1,9 @@
+"""
+Compliance module for aixtools.
+
+This module provides utilities for managing compliance-related data and operations.
+"""
+
+from .private_data import PrivateData
+
+__all__ = ["PrivateData"]

{aixtools-0.1.9 → aixtools-0.1.11}/aixtools/compliance/private_data.py

@@ -1,3 +1,5 @@
+"""Private data management module for aixtools compliance."""
+
 import json
 from pathlib import Path
 
@@ -90,7 +92,7 @@ class PrivateData:
 
     def _has_private_data_file(self) -> bool:
         """Check if the private data file exists in the workspace."""
-        private_data_path = self.get_private_data_path()  # type: ignore
+        private_data_path = self._get_private_data_path()
         return private_data_path.exists()
 
     def save(self) -> None:
@@ -102,7 +104,7 @@ class PrivateData:
             return
         # If there is private data, serialize this object as JSON
         private_data_path.parent.mkdir(parents=True, exist_ok=True)
-        with open(private_data_path, "w") as f:
+        with open(private_data_path, "w", encoding="utf-8") as f:
             # Dump class as JSON, excluding the context
             data_dict = self.__dict__.copy()
             data_dict["ctx"] = None
@@ -118,7 +120,7 @@ class PrivateData:
             self._private_datasets = []
             self._idap_datasets = []
             return
-        with open(private_data_path, "r") as f:
+        with open(private_data_path, "r", encoding="utf-8") as f:
             data = json.load(f)
             self.has_private_data = data.get("_has_private_data", False)
             self._private_datasets = data.get("_private_datasets", [])

{aixtools-0.1.9 → aixtools-0.1.11}/aixtools/mcp/client.py

@@ -1,12 +1,18 @@
 """MCP server utilities with caching and robust error handling."""
 
 import asyncio
-from typing import Any
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from typing import Any, AsyncGenerator
 
 import anyio
+import httpx
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from cachebox import TTLCache
 from mcp import types as mcp_types
+from mcp.client import streamable_http
 from mcp.shared.exceptions import McpError
+from mcp.shared.message import SessionMessage
 from pydantic_ai import RunContext, exceptions
 from pydantic_ai.mcp import MCPServerStreamableHTTP, ToolResult
 from pydantic_ai.toolsets.abstract import ToolsetTool
@@ -16,6 +22,7 @@ from aixtools.logging.logging_config import get_logger
 
 MCP_TOOL_CACHE_TTL = 300  # 5 minutes
 DEFAULT_MCP_CONNECTION_TIMEOUT = 30
+DEFAULT_MCP_READ_TIMEOUT = float(60 * 5)  # 5 minutes
 CACHE_KEY = "TOOL_LIST"
 
 logger = get_logger(__name__)
@@ -145,6 +152,23 @@ class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):
                 logger.warning("MCP %s: %s exception %s: %s", self.url, func.__name__, type(exc), exc)
             return fallback(exc)
 
+    @property
+    def _transport_client(self):
+        """Override base transport client with wrapper logging and suppressing exceptions"""
+        return patched_streamablehttp_client
+
+    @asynccontextmanager
+    async def client_streams(self):
+        """Override base client_streams with wrapper logging and suppressing exceptions"""
+        try:
+            async with super().client_streams() as streams:  # pylint: disable=contextmanager-generator-missing-cleanup
+                try:
+                    yield streams
+                except Exception as exc:  # pylint: disable=broad-except
+                    logger.error("MCP %s: client_streams; %s: %s", self.url, type(exc).__name__, exc)
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.error("MCP %s: client_streams: %s: %s", self.url, type(exc).__name__, exc)
+
     async def __aenter__(self):
         """Enter the context of the cached MCP server with complete cancellation isolation."""
         async with self._isolation_lock:
@@ -272,3 +296,80 @@ class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):
             raise exceptions.ModelRetry(text)
 
         return content[0] if len(content) == 1 else content
+
+
+class PatchedStreamableHTTPTransport(streamable_http.StreamableHTTPTransport):
+    """Patched StreamableHTTPTransport with exception suppression for _handle_post_request."""
+
+    async def _handle_post_request(self, ctx: streamable_http.RequestContext) -> None:
+        """Patched _handle_post_request with proper error handling."""
+        try:
+            await super()._handle_post_request(ctx)
+        except Exception as exc:  # pylint: disable=broad-except
+            logger.error("MCP %s: _handle_post_request %s: %s", self.url, type(exc).__name__, exc)
+
+
+@asynccontextmanager
+async def patched_streamablehttp_client(  # noqa: PLR0913, pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
+    url: str,
+    headers: dict[str, str] | None = None,
+    timeout: float | timedelta = 30,
+    sse_read_timeout: float | timedelta = DEFAULT_MCP_READ_TIMEOUT,
+    terminate_on_close: bool = True,
+    httpx_client_factory: streamable_http.McpHttpClientFactory = streamable_http.create_mcp_http_client,
+    auth: httpx.Auth | None = None,
+) -> AsyncGenerator[
+    tuple[
+        MemoryObjectReceiveStream[SessionMessage | Exception],
+        MemoryObjectSendStream[SessionMessage],
+        streamable_http.GetSessionIdCallback,
+    ],
+    None,
+]:
+    """Patched version of `streamablehttp_client` with exception suppression."""
+    try:
+        transport = PatchedStreamableHTTPTransport(url, headers, timeout, sse_read_timeout, auth)
+
+        read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
+        write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
+        async with anyio.create_task_group() as tg:
+            try:
+                async with httpx_client_factory(
+                    headers=transport.request_headers,
+                    timeout=httpx.Timeout(transport.timeout, read=transport.sse_read_timeout),
+                    auth=transport.auth,
+                ) as client:
+                    # Define callbacks that need access to tg
+                    def start_get_stream() -> None:
+                        tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
+
+                    tg.start_soon(
+                        transport.post_writer,
+                        client,
+                        write_stream_reader,
+                        read_stream_writer,
+                        write_stream,
+                        start_get_stream,
+                        tg,
+                    )
+
+                    try:
+                        yield (
+                            read_stream,
+                            write_stream,
+                            transport.get_session_id,
+                        )
+                    except GeneratorExit:
+                        logger.warning("patched_streamablehttp_client: GeneratorExit caught, closing streams.")
+                    finally:
+                        if transport.session_id and terminate_on_close:
+                            await transport.terminate_session(client)
+                        tg.cancel_scope.cancel()
+            finally:
+                await read_stream_writer.aclose()
+                await write_stream.aclose()
+    except Exception as exc:  # pylint: disable=broad-except
+        if str(exc) == "Attempted to exit cancel scope in a different task than it was entered in":
+            logger.warning("MCP %s: patched_streamablehttp_client: enter/exit cancel scope task mismatch.", url)
+        else:
+            logger.error("MCP %s: patched_streamablehttp_client: %s: %s", url, type(exc).__name__, exc)

{aixtools-0.1.9 → aixtools-0.1.11}/aixtools/testing/aix_test_model.py

@@ -108,6 +108,8 @@ class AixTestModel(Model):
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
+        *args,  # pylint: disable=unused-argument # Accept additional arguments for compatibility with pydantic-ai 1.0.9
+        **kwargs,  # pylint: disable=unused-argument
     ) -> AsyncIterator[StreamedResponse]:
         model_response = await self._request(messages, model_settings, model_request_parameters)
         yield TestStreamedResponse(_model_name=self.model_name, _structured_response=model_response, _messages=messages)

{aixtools-0.1.9 → aixtools-0.1.11}/aixtools.egg-info/SOURCES.txt

@@ -1,26 +1,10 @@
-.env_template
-.gitignore
-.python-version
 README.md
-README.ori.md
 pyproject.toml
-uv.lock
-.github/workflows/lint-and-test.yml
-.github/workflows/release.yml
-.roo/rules/rules-mcp.md
-.roo/rules/rules.md
-.vscode/settings.json
 aixtools/__init__.py
 aixtools/_version.py
 aixtools/app.py
 aixtools/chainlit.md
 aixtools/context.py
-aixtools.egg-info/PKG-INFO
-aixtools.egg-info/SOURCES.txt
-aixtools.egg-info/dependency_links.txt
-aixtools.egg-info/entry_points.txt
-aixtools.egg-info/requires.txt
-aixtools.egg-info/top_level.txt
 aixtools/.chainlit/config.toml
 aixtools/.chainlit/translations/bn.json
 aixtools/.chainlit/translations/en-US.json
@@ -47,6 +31,7 @@ aixtools/agents/__init__.py
 aixtools/agents/agent.py
 aixtools/agents/agent_batch.py
 aixtools/agents/prompt.py
+aixtools/compliance/__init__.py
 aixtools/compliance/private_data.py
 aixtools/db/__init__.py
 aixtools/db/database.py
@@ -97,41 +82,4 @@ aixtools/utils/utils.py
 aixtools/utils/chainlit/cl_agent_show.py
 aixtools/utils/chainlit/cl_utils.py
 aixtools/vault/__init__.py
-aixtools/vault/vault.py
-docker/mcp-base/Dockerfile
-docker/mcp-base/zscaler.crt
-notebooks/example_faulty_mcp_server.ipynb
-notebooks/example_mcp_server_stdio.ipynb
-notebooks/example_raw_mcp_client.ipynb
-notebooks/example_tool_doctor.ipynb
-scripts/config.sh
-scripts/lint.sh
-scripts/log_view.sh
-scripts/run_example_mcp_server.sh
-scripts/run_faulty_mcp_server.sh
-scripts/run_server.sh
-scripts/test.sh
-tests/__init__.py
-tests/unit/__init__.py
-tests/unit/a2a/__init__.py
-tests/unit/a2a/google_sdk/__init__.py
-tests/unit/a2a/google_sdk/test_card.py
-tests/unit/a2a/google_sdk/test_remote_agent_connection.py
-tests/unit/a2a/google_sdk/test_utils.py
-tests/unit/a2a/google_sdk/pydantic_ai_adapter/__init__.py
-tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_agent_executor.py
-tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_storage.py
-tests/unit/agents/__init__.py
-tests/unit/agents/test_prompt.py
-tests/unit/compliance/test_private_data.py
-tests/unit/google/__init__.py
-tests/unit/google/test_client.py
-tests/unit/mcp/__init__.py
-tests/unit/mcp/test_client.py
-tests/unit/server/__init__.py
-tests/unit/server/test_path.py
-tests/unit/server/test_utils.py
-tests/unit/utils/__init__.py
-tests/unit/utils/test_files.py
-tests/unit/vault/__init__.py
-tests/unit/vault/test_vault.py
+aixtools/vault/vault.py

{aixtools-0.1.9 → aixtools-0.1.11}/pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "langchain-ollama>=0.3.2",
     "langchain-openai>=0.3.14",
     "mcp>=1.11.0",
+    "mypy>=1.18.2",
     "pandas>=2.2.3",
     "pydantic-ai>=0.4.10",
     "pylint>=3.3.7",

aixtools-0.1.9/.env_template

@@ -1,49 +0,0 @@
-#-----------------------------------------------------------------------------
-#
-# Template for .env file
-#
-#-----------------------------------------------------------------------------
-
-# Select which credentials to use
-# Options: azure, openai, bedrock, ollama, openrouter
-MODEL_FAMILY=azure
-VDB_EMBEDDINGS_MODEL_FAMILY=azure
-
-MODEL_TIMEOUT=120
-
-# Azure OpenAI credentials
-# If 'MODEL_FAMILY=azure'
-AZURE_OPENAI_ENDPOINT=https://your_endpoint.openai.azure.com
-AZURE_OPENAI_API_VERSION=2024-06-01
-AZURE_OPENAI_API_KEY=your_secret_key
-AZURE_MODEL_NAME=gpt-4o
-AZURE_OPENAI_PROVIDER_ID=azure
-AZURE_VDB_EMBEDDINGS_MODEL_NAME=text-embedding-3-small
-
-# OpenAI credentials
-# If 'MODEL_FAMILY=openai'
-OPENAI_MODEL_NAME=gpt-4.5-preview
-OPENAI_API_KEY=openai_api_key
-
-# Bedrock credentials
-# If 'MODEL_FAMILY=bedrock'
-BEDROCK_MODEL_NAME=us.anthropic.claude-sonnet-4-20250514-v1:0 # or any other enabled model
-AWS_REGION=us-east-1
-
-# OpenRouter credentials
-# If 'MODEL_FAMILY=openrouter'
-OPENROUTER_API_URL=https://openrouter.ai/api/v1
-OPENROUTER_MODEL_NAME=anthropic/claude-3.7-sonnet
-OPENROUTER_API_KEY=your_openrouter_api_key
-
-# Ollama models
-# If 'MODEL_FAMILY=ollama'
-OLLAMA_MODEL_NAME=llama3.2:3b-instruct-fp16
-OLLAMA_URL=http://localhost:11434/v1
-
-# Vector DB embeddings
-OPENAI_VDB_EMBEDDINGS_MODEL_NAME=text-embedding-3-small
-OLLAMA_VDB_EMBEDDINGS_MODEL_NAME=snowflake-arctic-embed2:latest
-
-# Custom certificate for Docker builds (optional)
-# CUSTOM_CERT_FILE=/path/to/your/certificate.crt

aixtools-0.1.9/.github/workflows/lint-and-test.yml

@@ -1,41 +0,0 @@
-name: Lint and Test
-
-on:
-  push:
-    branches: [main]
-  pull_request:
-    branches: [main]
-
-jobs:
-  lint-and-test:
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python 3.12
-        uses: actions/setup-python@v4
-        with:
-          python-version: "3.12"
-
-      - name: Install uv
-        run: |
-          pip install uv
-
-      - name: Install dependencies
-        run: |
-          uv sync
-          uv pip install -e ".[feature]"
-
-      - name: Set up environment
-        run: |
-          touch .env
-
-      - name: Run linting
-        run: |
-          ./scripts/lint.sh
-
-      - name: Test
-        run: |
-          ./scripts/test.sh

aixtools-0.1.9/.github/workflows/release.yml

@@ -1,58 +0,0 @@
-name: Release
-
-on:
-  release:
-    types: [published]
-
-jobs:
-  build-and-publish:
-    runs-on: ubuntu-latest
-    environment: release
-    env:
-      DOCKER_IMAGE_NAME: mcp-base
-      AZ_ARTIFACTORY_ODSP_DOCKER_REGISTRY: docker-l-odsp.artifactory.astrazeneca.com
-      AZ_ARTIFACTORY_ODSP_DOCKER_NAMESPACE: navari
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0 # Fetch full history for setuptools_scm
-
-      - name: Install uv
-        uses: astral-sh/setup-uv@v4
-
-      - name: Install build dependencies
-        run: uv sync --group dev
-
-      - name: Build Python package
-        run: uv run python -m build
-
-      - name: Check Python package contents
-        run: |
-          ls -la dist/
-          uv run python -m twine check dist/*
-
-      - name: Publish Python package to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          password: ${{ secrets.PYPI_API_TOKEN }}
-          skip-existing: true
-
-      - name: Login to Docker Registry (JFrog Artifactory 🐸)
-        uses: docker/login-action@v3
-        with:
-          registry: ${{ env.AZ_ARTIFACTORY_ODSP_DOCKER_REGISTRY }}
-          username: ${{ vars.AZ_ARTIFACTORY_ODSP_DOCKER_USER }}
-          password: ${{ secrets.AZ_ARTIFACTORY_ODSP_DOCKER_PASSWORD }}
-
-      - name: Build and push Docker image 🐳
-        id: docker_build
-        uses: docker/build-push-action@v6
-        env:
-          AZ_ARTIFACTORY_ODSP_DOCKER_IMAGE_REF: ${{ env.AZ_ARTIFACTORY_ODSP_DOCKER_REGISTRY }}/${{ env.AZ_ARTIFACTORY_ODSP_DOCKER_NAMESPACE }}/${{ env.DOCKER_IMAGE_NAME }}
-        with:
-          push: true
-          context: docker/${{ env.DOCKER_IMAGE_NAME }}
-          tags: |
-            ${{ env.AZ_ARTIFACTORY_ODSP_DOCKER_IMAGE_REF }}:${{ github.ref_name }}

aixtools-0.1.9/.gitignore

@@ -1,188 +0,0 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-.DS_Store
-agents/.DS_Store
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# UV
-#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#uv.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
-.pdm.toml
-.pdm-python
-.pdm-build/
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-
-#IDEs
-.idea/
-
-# Ruff stuff:
-.ruff_cache/
-
-# PyPI configuration file
-.pypirc
-
-# Chainlit stuff
-.files/
-
-# Data directories
-data/
-aixtools/data/
-logs/
-aixtools/_version.py

aixtools-0.1.9/.python-version

@@ -1 +0,0 @@
-3.12.2