aixtools 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.5'
-__version_tuple__ = version_tuple = (0, 1, 5)
+__version__ = version = '0.1.6'
+__version_tuple__ = version_tuple = (0, 1, 6)
 
 __commit_id__ = commit_id = None

aixtools/utils/config.py

@@ -122,3 +122,10 @@ LOGFIRE_TRACES_ENDPOINT = get_variable_env("LOGFIRE_TRACES_ENDPOINT", True, "")
 GOOGLE_GENAI_USE_VERTEXAI = str2bool(get_variable_env("GOOGLE_GENAI_USE_VERTEXAI", True, True))
 GOOGLE_CLOUD_PROJECT = get_variable_env("GOOGLE_CLOUD_PROJECT", True)
 GOOGLE_CLOUD_LOCATION = get_variable_env("GOOGLE_CLOUD_LOCATION", True)
+
+# vault parameters.
+VAULT_ADDRESS = get_variable_env("VAULT_ADDRESS", default="http://localhost:8200")
+VAULT_TOKEN = get_variable_env("VAULT_TOKEN", default="vault-token")
+VAULT_ENV = get_variable_env("ENV", default="dev")
+VAULT_MOUNT_POINT = get_variable_env("VAULT_MOUNT_POINT", default="secret")
+VAULT_PATH_PREFIX = get_variable_env("VAULT_PATH_PREFIX", default="path")

aixtools/vault/__init__.py

@@ -0,0 +1,7 @@
+"""
+Provides a Vault client for storing and retrieving user service api keys.
+"""
+
+from .vault import VaultClient
+
+__all__ = ["VaultClient"]

aixtools/vault/vault.py

@@ -0,0 +1,73 @@
+# ruff: noqa: PLR0913
+"""
+Provides a Vault client for storing and retrieving user service api keys.
+"""
+
+import logging
+from typing import Optional
+
+import hvac
+from hvac.exceptions import InvalidPath
+
+from aixtools.utils import config
+
+logger = logging.getLogger(__name__)
+
+
+class VaultAuthError(Exception):
+    """Exception raised for vault authentication errors."""
+
+
+class VaultClient:
+    """Vault client for storing and retrieving user service api keys."""
+
+    def __init__(self):
+        self.client = hvac.Client(url=config.VAULT_ADDRESS, token=config.VAULT_TOKEN)
+
+        if not self.client.is_authenticated():
+            raise VaultAuthError("Vault client authentication failed. Check vault_token.")
+
+    def store_user_service_api_key(self, *, user_id: str, service_name: str, user_api_key: str):
+        """
+        Store user's service api key in the Vault at the specified vault mount
+        point, where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+        """
+        secret_path = None
+        try:
+            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+            print("secret_path", secret_path)
+            secret_dict = {"user-api-key": user_api_key}
+            self.client.secrets.kv.v2.create_or_update_secret(
+                secret_path, secret=secret_dict, mount_point=config.VAULT_MOUNT_POINT
+            )
+
+            logger.info("Secret written to path %s", secret_path)
+        except Exception as e:
+            logger.error("Failed to write secret to path %s: %s", secret_path, str(e))
+            raise VaultAuthError(e) from e
+
+    def read_user_service_api_key(self, *, user_id: str, service_name) -> Optional[str]:
+        """
+        Read user's service api key in from vault at the specified mount point,
+        where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+        """
+        secret_path = None
+
+        try:
+            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+            logger.info("Reading secret from path %s", secret_path)
+            response = self.client.secrets.kv.v2.read_secret_version(
+                secret_path, mount_point=config.VAULT_MOUNT_POINT, raise_on_deleted_version=True
+            )
+            secret_data = response["data"]["data"]
+            user_api_key = secret_data["user-api-key"]
+            logger.info("Secret read from path %s ", secret_path)
+            return user_api_key
+        except InvalidPath:
+            # Secret path does not exist
+            logger.warning("Secret path does not exist %s ", secret_path)
+            return None
+
+        except Exception as e:
+            logger.error("Failed to read secret from path %s: %s", secret_path, str(e))
+            raise VaultAuthError(e) from e

aixtools-0.1.6.dist-info/METADATA

@@ -0,0 +1,668 @@
+Metadata-Version: 2.4
+Name: aixtools
+Version: 0.1.6
+Summary: Tools for AI exploration and debugging
+Requires-Python: >=3.11.2
+Description-Content-Type: text/markdown
+Requires-Dist: a2a-sdk>=0.3.1
+Requires-Dist: cachebox>=5.0.1
+Requires-Dist: chainlit>=2.5.5
+Requires-Dist: colorlog>=6.9.0
+Requires-Dist: fasta2a>=0.5.0
+Requires-Dist: fastmcp>=2.10.2
+Requires-Dist: hvac>=2.3.0
+Requires-Dist: ipykernel>=6.29.5
+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: pandas>=2.2.3
+Requires-Dist: pydantic-ai>=0.4.10
+Requires-Dist: pylint>=3.3.7
+Requires-Dist: rich>=14.0.0
+Requires-Dist: ruff>=0.11.6
+Requires-Dist: streamlit>=1.44.1
+Requires-Dist: watchdog>=6.0.0
+Provides-Extra: test
+Requires-Dist: pyyaml; extra == "test"
+Provides-Extra: feature
+Requires-Dist: logfire; extra == "feature"
+
+# AIXtools
+
+AIXtools is a comprehensive Python library for AI agent development, debugging, and deployment. It provides a complete toolkit for building, testing, and monitoring AI agents with support for multiple model providers, advanced logging, and agent-to-agent communication.
+
+## Capabilities
+
+Agents
+- Agent Development & Management - `aixtools/agents/`
+- Agent Batch Processing - `aixtools/agents/agent_batch.py`
+- Agent Prompting System - `aixtools/agents/prompt.py`
+
+A2A
+- Agent-to-Agent Communication (A2A) - `aixtools/a2a/`
+- Google SDK Integration for A2A - `aixtools/a2a/google_sdk/`
+- PydanticAI Adapter for Google SDK - `aixtools/a2a/google_sdk/pydantic_ai_adapter/`
+
+Databases
+- Database Integration - `aixtools/db/`
+- Vector Database Support - `aixtools/db/vector_db.py`
+
+Logging & Debugging
+- Log Viewing Application - `aixtools/log_view/`
+- Object Logging System - `aixtools/logging/`
+- Model Patch Logging - `aixtools/logging/model_patch_logging.py`
+- Log Filtering System - `aixtools/logfilters/`
+- FastMCP Logging - `aixtools/mcp/fast_mcp_log.py`
+- Command Line Interface for Log Viewing - Entry point: `log_view`
+- MCP (Model Context Protocol) Support - `aixtools/logging/mcp_log_models.py`, `aixtools/logging/mcp_logger.py`
+
+Testing & Tools
+- Testing Utilities - `aixtools/testing/`
+- Mock Tool System - `aixtools/testing/mock_tool.py`
+- Model Patch Caching - `aixtools/testing/model_patch_cache.py`
+- Tool Doctor System - `aixtools/tools/doctor/`
+- Tool Recommendation Engine - `aixtools/tools/doctor/tool_recommendation.py`
+- FaultyMCP - `aixtools/mcp/faulty_mcp.py`
+
+Chainlit & HTTP Server
+- Chainlit Integration - `aixtools/app.py`, `aixtools/chainlit.md`
+- Chainlit Utilities - `aixtools/utils/chainlit/`
+- HTTP Server Framework - `aixtools/server/`
+- App Mounting System - `aixtools/server/app_mounter.py`
+
+Programming utils
+- Persisted Dictionary - `aixtools/utils/persisted_dict.py`
+- Enum with Description - `aixtools/utils/enum_with_description.py`
+- Context Management - `aixtools/context.py`
+- Configuration Management - `aixtools/utils/config.py`, `aixtools/utils/config_util.py`
+- File Utilities - `aixtools/utils/files.py`
+
+## Installation
+
+### From GitHub
+
+```bash
+uv add aixtools
+```
+
+### Development Setup
+
+```bash
+# Create a new project
+uv init MyNewProject
+cd MyNewProject
+
+# Add virtual environment and activate it
+uv venv .venv
+source .venv/bin/activate
+
+# Add this package
+uv add aixtools
+```
+
+### Updating
+
+```bash
+uv add --upgrade aixtools
+```
+
+## Environment Configuration
+
+AIXtools requires environment variables for model providers. 
+
+**IMPORTANT:** Create a `.env` file based on [`.env_template`](./.env_template):
+
+```bash
+# Model family (azure, openai, or ollama)
+MODEL_FAMILY=azure
+MODEL_TIMEOUT=120
+
+# Azure OpenAI
+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
+
+# OpenAI
+OPENAI_MODEL_NAME=gpt-4.5-preview
+OPENAI_API_KEY=openai_api_key
+
+# Ollama
+OLLAMA_MODEL_NAME=llama3.2:3b-instruct-fp16
+OLLAMA_LOCAL_URL=http://localhost:11434/v1
+```
+
+## Agents
+
+### Basic Agent Usage
+
+```python
+from aixtools.agents.agent import get_agent, run_agent
+
+async def main():
+    agent = get_agent(system_prompt="You are a helpful assistant.")
+    result, nodes = await run_agent(agent, "Explain quantum computing")
+    print(result)
+```
+
+### Agent Development & Management
+
+The agent system provides a unified interface for creating and managing AI agents across different model providers.
+
+```python
+from aixtools.agents.agent import get_agent, run_agent
+
+# Create an agent with default model
+agent = get_agent(system_prompt="You are a helpful assistant.")
+
+# Run the agent
+result, nodes = await run_agent(agent, "Tell me about AI")
+```
+
+### Agent Batch Processing
+
+Process multiple agent queries simultaneously with built-in concurrency control and result aggregation.
+
+```python
+from aixtools.agents.agent_batch import agent_batch, AgentQueryParams
+
+# Create query parameters
+query_parameters = [
+    AgentQueryParams(prompt="What is the meaning of life"),
+    AgentQueryParams(prompt="Who is the prime minister of Canada")
+]
+
+# Run queries in batches
+async for result in agent_batch(query_parameters):
+    print(result)
+```
+
+## A2A (Agent-to-Agent Communication)
+
+The A2A module provides a comprehensive framework for enabling sophisticated communication between AI agents across different environments and platforms. It includes Google SDK integration, PydanticAI adapters, and FastA2A application conversion capabilities.
+
+### Core Features
+
+**Agent Application Conversion**
+- Convert PydanticAI agents into FastA2A applications
+- Support for session metadata extraction and context management
+- Custom worker classes with enhanced data part support
+- Automatic handling of user and session identification
+
+**Remote Agent Connections**
+- Establish connections between agents across different environments
+- Asynchronous message sending with task polling capabilities
+- Terminal state detection and error handling
+- Support for various message types including text, files, and data
+
+**Google SDK Integration**
+- Native integration with Google's A2A SDK
+- Card-based agent representation and discovery
+- PydanticAI adapter for seamless Google SDK compatibility
+- Storage and execution management for agent interactions
+
+### Agent-to-Agent Communication (A2A)
+
+Enable sophisticated agent interactions with Google SDK integration and PydanticAI adapters.
+
+```python
+from aixtools.a2a.google_sdk.remote_agent_connection import RemoteAgentConnection
+from aixtools.a2a.app import agent_to_a2a
+
+# Convert a PydanticAI agent to FastA2A application
+a2a_app = agent_to_a2a(
+    agent=my_agent,
+    name="MyAgent",
+    description="A helpful AI assistant",
+    skills=[{"name": "chat", "description": "General conversation"}]
+)
+
+# Connect agents across different environments
+connection = RemoteAgentConnection(card=agent_card, client=a2a_client)
+response = await connection.send_message_with_polling(message)
+```
+
+## Databases
+
+### Database Integration
+
+Support for both traditional and vector databases with seamless integration.
+
+```python
+from aixtools.db.database import Database
+from aixtools.db.vector_db import VectorDB
+
+# Traditional database
+db = Database("sqlite:///app.db")
+
+# Vector database for embeddings
+vector_db = VectorDB()
+vector_db.add_documents(documents)
+```
+
+## Logging & Debugging
+
+AixTools provides functionality for logging and debugging.
+
+### Basic Logging and Debugging
+
+```python
+from aixtools.agents.agent import get_agent, run_agent
+
+async def main():
+    # Create an agent
+    agent = get_agent(system_prompt="You are a helpful assistant.")
+    
+    # Run agent - logging is automatic via ObjectLogger
+    result, nodes = await run_agent(
+        agent,
+        "Explain quantum computing",
+        debug=True,  # Enable debug logging
+        log_model_requests=True  # Log model requests/responses
+    )
+    
+    print(f"Result: {result}")
+    print(f"Logged {len(nodes)} nodes")
+```
+
+### Log Viewing Application
+
+Interactive Streamlit application for analyzing logged objects and debugging agent behavior.
+
+**Features:**
+- Log file selection and filtering
+- Node visualization with expand/collapse
+- Export capabilities to JSON
+- Regex pattern matching
+- Real-time log monitoring
+
+```bash
+# Run the log viewer
+log_view
+
+# Or specify custom log directory
+log_view /path/to/logs
+```
+
+### Object Logging & Debugging
+
+Advanced logging system with object serialization and visual debugging tools.
+
+```python
+from aixtools.logging.log_objects import ObjectLogger
+
+# Log any pickleable object
+with ObjectLogger() as logger:
+    logger.log({"message": "Hello, world!"})
+    logger.log(agent_response)
+```
+
+### MCP Logger
+
+This is an MCP server that can log MCP requests and responses.
+
+```python
+from aixtools.mcp.fast_mcp_log import FastMcpLog
+
+# Use FastMCP server with logging
+mcp = FastMcpLog("Demo")
+```
+
+### Model Patching System
+
+Dynamic model behavior modification for testing and debugging.
+
+```python
+from aixtools.model_patch.model_patch import ModelPatch
+
+# Apply patches to models for testing
+with ModelPatch() as patch:
+    patch.apply_response_override("test response")
+    result = await agent.run("test prompt")
+```
+
+### FaultyMCP
+
+A specialized MCP server designed for testing error handling and resilience in MCP client implementations. FaultyMCP simulates various failure scenarios including network errors, server crashes, and random exceptions.
+
+**Features:**
+- Configurable error probabilities for different request types
+- HTTP 404 error injection for POST/DELETE requests
+- Server crash simulation on GET requests
+- Random exception throwing in tool operations
+- MCP-specific error simulation (ValidationError, ResourceError, etc.)
+- Safe mode for controlled testing
+
+```python
+from aixtools.mcp.faulty_mcp import run_server_on_port, config
+
+# Configure error probabilities
+config.prob_on_post_404 = 0.3      # 30% chance of 404 on POST
+config.prob_on_get_crash = 0.1     # 10% chance of crash on GET
+config.prob_in_list_tools_throw = 0.2  # 20% chance of exception in tools/list
+
+# Run the faulty server
+run_server_on_port()
+```
+
+**Command Line Usage:**
+```bash
+# Run with default error probabilities
+python -m aixtools.mcp.faulty_mcp
+
+# Run in safe mode (no errors by default)
+python -m aixtools.mcp.faulty_mcp --safe-mode
+
+# Custom configuration
+python -m aixtools.mcp.faulty_mcp \
+    --port 8888 \
+    --prob-on-post-404 0.2 \
+    --prob-on-get-crash 0.1 \
+    --prob-in-list-tools-throw 0.3
+```
+
+By default, the "FaultyMCP" includes several tools you can use in your tests:
+- `add(a, b)` - Basic addition (reliable)
+- `multiply(a, b)` - Basic multiplication (reliable)
+- `always_error()` - Always throws an exception
+- `random_throw_exception(a, b, prob)` - Randomly throws exceptions
+- `freeze_server(seconds)` - Simulates server freeze
+- `throw_404_exception()` - Throws HTTP 404 error
+
+## Testing & Tools
+
+AIXtools provides comprehensive testing utilities and diagnostic tools for AI agent development and debugging.
+
+### Testing Utilities
+
+The testing module provides mock tools, model patching, and test utilities for comprehensive agent testing.
+
+```python
+from aixtools.testing.mock_tool import MockTool
+from aixtools.testing.model_patch_cache import ModelPatchCache
+from aixtools.testing.aix_test_model import AixTestModel
+
+# Create mock tools for testing
+mock_tool = MockTool(name="test_tool", response="mock response")
+
+# Use model patch caching for consistent test results
+cache = ModelPatchCache()
+cached_response = cache.get_cached_response("test_prompt")
+
+# Test model for controlled testing scenarios
+test_model = AixTestModel()
+```
+
+### Tool Doctor System
+
+Automated tool analysis and recommendation system for optimizing agent tool usage.
+
+```python
+from aixtools.tools.doctor.tool_doctor import ToolDoctor
+from aixtools.tools.doctor.tool_recommendation import ToolRecommendation
+
+# Analyze tool usage patterns
+doctor = ToolDoctor()
+analysis = doctor.analyze_tools(agent_logs)
+
+# Get tool recommendations
+recommendation = ToolRecommendation()
+suggestions = recommendation.recommend_tools(agent_context)
+```
+
+### Mock Tool System
+
+Create and manage mock tools for testing agent behavior without external dependencies.
+
+```python
+from aixtools.testing.mock_tool import MockTool
+
+# Create a mock tool with predefined responses
+mock_calculator = MockTool(
+    name="calculator",
+    description="Performs mathematical calculations",
+    response_map={
+        "2+2": "4",
+        "10*5": "50"
+    }
+)
+
+# Use in agent testing
+agent = get_agent(tools=[mock_calculator])
+result = await run_agent(agent, "What is 2+2?")
+```
+
+### Model Patch Caching
+
+Cache model responses for consistent testing and development workflows.
+
+```python
+from aixtools.testing.model_patch_cache import ModelPatchCache
+
+# Initialize cache
+cache = ModelPatchCache(cache_dir="./test_cache")
+
+# Cache responses for specific prompts
+cache.cache_response("test prompt", "cached response")
+
+# Retrieve cached responses
+response = cache.get_cached_response("test prompt")
+```
+
+### FaultyMCP Testing Server
+
+Specialized MCP server for testing error handling and resilience in MCP implementations.
+
+```python
+from aixtools.mcp.faulty_mcp import run_server_on_port, config
+
+# Configure error probabilities for testing
+config.prob_on_post_404 = 0.3      # 30% chance of 404 on POST
+config.prob_on_get_crash = 0.1     # 10% chance of crash on GET
+config.prob_in_list_tools_throw = 0.2  # 20% chance of exception
+
+# Run the faulty server for testing
+run_server_on_port(port=8888)
+```
+
+**Available Test Tools:**
+- `add(a, b)` - Reliable addition operation
+- `multiply(a, b)` - Reliable multiplication operation
+- `always_error()` - Always throws an exception
+- `random_throw_exception(a, b, prob)` - Randomly throws exceptions
+- `freeze_server(seconds)` - Simulates server freeze
+- `throw_404_exception()` - Throws HTTP 404 error
+
+**Command Line Usage:**
+```bash
+# Run with default error probabilities
+python -m aixtools.mcp.faulty_mcp
+
+# Run in safe mode (no errors)
+python -m aixtools.mcp.faulty_mcp --safe-mode
+
+# Custom configuration
+python -m aixtools.mcp.faulty_mcp \
+    --port 8888 \
+    --prob-on-post-404 0.2 \
+    --prob-on-get-crash 0.1
+```
+
+### Running Tests
+
+Execute the test suite using the provided scripts:
+
+```bash
+# Run all tests
+./scripts/test.sh
+
+# Run unit tests only
+./scripts/test_unit.sh
+
+# Run integration tests only
+./scripts/test_integration.sh
+```
+
+## Chainlit & HTTP Server
+
+### Chainlit Integration
+
+Ready-to-use Chainlit application for interactive agent interfaces.
+
+```python
+# Run the Chainlit app
+# Configuration in aixtools/chainlit.md
+# Main app in aixtools/app.py
+```
+
+## Programming Utils
+
+AIXtools provides essential programming utilities for configuration management, data persistence, file operations, and context handling.
+
+### Persisted Dictionary
+
+Persistent key-value storage with automatic serialization and file-based persistence.
+
+```python
+from aixtools.utils.persisted_dict import PersistedDict
+
+# Create a persistent dictionary
+cache = PersistedDict("cache.json")
+
+# Store and retrieve data
+cache["user_preferences"] = {"theme": "dark", "language": "en"}
+cache["session_data"] = {"last_login": "2024-01-01"}
+
+# Data is automatically saved to file
+print(cache["user_preferences"])  # Persists across program restarts
+```
+
+### Enum with Description
+
+Enhanced enum classes with built-in descriptions for better documentation and user interfaces.
+
+```python
+from aixtools.utils.enum_with_description import EnumWithDescription
+
+class ModelType(EnumWithDescription):
+    GPT4 = ("gpt-4", "OpenAI GPT-4 model")
+    CLAUDE = ("claude-3", "Anthropic Claude-3 model")
+    LLAMA = ("llama-2", "Meta LLaMA-2 model")
+
+# Access enum values and descriptions
+print(ModelType.GPT4.value)        # "gpt-4"
+print(ModelType.GPT4.description)  # "OpenAI GPT-4 model"
+
+# Get all descriptions
+for model in ModelType:
+    print(f"{model.value}: {model.description}")
+```
+
+### Context Management
+
+Centralized context management for sharing state across components.
+
+```python
+from aixtools.context import Context
+
+# Create and use context
+context = Context()
+context.set("user_id", "12345")
+context.set("session_data", {"preferences": {"theme": "dark"}})
+
+# Retrieve context data
+user_id = context.get("user_id")
+session_data = context.get("session_data")
+
+# Context can be passed between components
+def process_request(ctx: Context):
+    user_id = ctx.get("user_id")
+    # Process with user context
+```
+
+### Configuration Management
+
+Robust configuration handling with environment variable support and validation.
+
+```python
+from aixtools.utils.config import Config
+from aixtools.utils.config_util import load_config
+
+# Load configuration from environment and files
+config = load_config()
+
+# Access configuration values
+model_name = config.get("MODEL_NAME", "gpt-4")
+api_key = config.get("API_KEY")
+timeout = config.get("TIMEOUT", 30, int)
+
+# Configuration with validation
+class AppConfig(Config):
+    model_name: str = "gpt-4"
+    max_tokens: int = 1000
+    temperature: float = 0.7
+
+app_config = AppConfig()
+```
+
+### File Utilities
+
+Enhanced file operations with Path support and utility functions.
+
+```python
+from aixtools.utils.files import read_file, write_file, ensure_directory
+from pathlib import Path
+
+# Read and write files with automatic encoding handling
+content = read_file("data.txt")
+write_file("output.txt", "Hello, world!")
+
+# Ensure directories exist
+data_dir = Path("data/logs")
+ensure_directory(data_dir)
+
+# Work with file paths
+config_path = Path("config") / "settings.json"
+if config_path.exists():
+    config_data = read_file(config_path)
+```
+
+### Chainlit Utilities
+
+Specialized utilities for Chainlit integration and agent display.
+
+```python
+from aixtools.utils.chainlit.cl_agent_show import show_agent_response
+from aixtools.utils.chainlit.cl_utils import format_message
+
+# Display agent responses in Chainlit
+await show_agent_response(
+    response="Hello, how can I help you?",
+    metadata={"model": "gpt-4", "tokens": 150}
+)
+
+# Format messages for Chainlit display
+formatted_msg = format_message(
+    content="Processing your request...",
+    message_type="info"
+)
+```
+
+### General Utilities
+
+Common utility functions for everyday programming tasks.
+
+```python
+from aixtools.utils.utils import safe_json_loads, timestamp_now, hash_string
+
+# Safe JSON parsing
+data = safe_json_loads('{"key": "value"}', default={})
+
+# Get current timestamp
+now = timestamp_now()
+
+# Generate hash for strings
+file_hash = hash_string("content to hash")
+```
+

{aixtools-0.1.5.dist-info → aixtools-0.1.6.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=rdxBMYpwzYxiWk08QbPLHSAxHoDfeKWwyaJIAM0lSic,704
+aixtools/_version.py,sha256=riGXiVTWXmtdoju9hVCWvTxpszEMAAIK0sZZWoLKlnU,704
 aixtools/app.py,sha256=JzQ0nrv_bjDQokllIlGHOV0HEb-V8N6k_nGQH-TEsVU,5227
 aixtools/chainlit.md,sha256=yC37Ly57vjKyiIvK4oUvf4DYxZCwH7iocTlx7bLeGLU,761
 aixtools/context.py,sha256=I_MD40ZnvRm5WPKAKqBUAdXIf8YaurkYUUHSVVy-QvU,598
@@ -68,7 +68,7 @@ aixtools/tools/doctor/__init__.py,sha256=FPwYzC1eJyw8IH0-BP0wgxSprLy6Y_4yXCek749
 aixtools/tools/doctor/tool_doctor.py,sha256=flp00mbFwVI0-Ii_xC4YDW6Vrn-EAExA1TtQkY6cOZE,2583
 aixtools/tools/doctor/tool_recommendation.py,sha256=t-l5bm6kwnXs1NH-ZZVTWhVrEAmWa460M44bi_Bip4g,1463
 aixtools/utils/__init__.py,sha256=xT6almZBQYMfj4h7Hq9QXDHyVXbOOTxqLsmJsxYYnSw,757
-aixtools/utils/config.py,sha256=x6zfDDAwD0dQglRJnlENxYXZKo_IMUaSLo9Ed6XKboc,4446
+aixtools/utils/config.py,sha256=pNWVMC1V9Hn2KEqaXaLbhxCI_iwQfGiVKKaZLrRM4ug,4820
 aixtools/utils/config_util.py,sha256=3Ya4Qqhj1RJ1qtTTykQ6iayf5uxlpigPXgEJlTi1wn4,2229
 aixtools/utils/enum_with_description.py,sha256=zjSzWxG74eR4x7dpmb74pLTYCWNSMvauHd7_9LpDYIw,1088
 aixtools/utils/files.py,sha256=8JnxwHJRJcjWCdFpjzWmo0po2fRg8esj4H7sOxElYXU,517
@@ -76,6 +76,8 @@ aixtools/utils/persisted_dict.py,sha256=0jQzV7oF-A6Or-HjcU6V7aMXWQL67SOKpULgmtFw
 aixtools/utils/utils.py,sha256=5911Ej1ES2NU_FKIWA3CWKhKnwgjvi1aDR2aiD6Xv3E,4880
 aixtools/utils/chainlit/cl_agent_show.py,sha256=vaRuowp4BRvhxEr5hw0zHEJ7iaSF_5bo_9BH7pGPPpw,4398
 aixtools/utils/chainlit/cl_utils.py,sha256=fxaxdkcZg6uHdM8uztxdPowg3a2f7VR7B26VPY4t-3c,5738
+aixtools/vault/__init__.py,sha256=fsr_NuX3GZ9WZ7dGfe0gp_5-z3URxAfwVRXw7Xyc0dU,141
+aixtools/vault/vault.py,sha256=JeAnRnsXm_etkTUntE9ABIiNlI8M81E2b85IQHfGehI,2817
 docker/mcp-base/Dockerfile,sha256=sSpbt0sasSBHHeGwPIpJpiEQMU5HGeXzerK8biVSt7Q,1547
 notebooks/example_faulty_mcp_server.ipynb,sha256=b2Cy3GXfj-gOBZ7SoUzj25F1rxp5u-32EWPHWQ-sxn8,1729
 notebooks/example_mcp_server_stdio.ipynb,sha256=ya4dRKNFU2vQxob-uIhKHGAzINXGQ6MehgKVmSCpHLk,1634
@@ -109,8 +111,10 @@ tests/unit/server/test_path.py,sha256=1QKiKLLRga9GNxmaUEt_wEZ9U14yzB-7PIhAOgB4ww
 tests/unit/server/test_utils.py,sha256=kvhzdgNfsJl5tqcRBWg2yTR5GPpyrFCOmEIOuHb3904,14848
 tests/unit/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/utils/test_files.py,sha256=AKFmXQqXstyKd2PreE4EmQyhQYeqOmu1Sp80MwHrf_Q,5782
-aixtools-0.1.5.dist-info/METADATA,sha256=QbJnhcWkwh9yMwKh5e95XmVuLKVAltsOx2bce11oKbs,10170
-aixtools-0.1.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.1.5.dist-info/entry_points.txt,sha256=dHoutULEZx7xXSqJrZdViSVjfInJibfLibi2nRXL3SE,56
-aixtools-0.1.5.dist-info/top_level.txt,sha256=ee4eF-0pqu45zCUVml0mWIhnXQgqMQper2-49BBVHLY,40
-aixtools-0.1.5.dist-info/RECORD,,
+tests/unit/vault/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/unit/vault/test_vault.py,sha256=YzmJ-JZo8OLGNcmc845knCh3zp8jhkANVgKT09cuhSU,4250
+aixtools-0.1.6.dist-info/METADATA,sha256=6JFv550ISVltjdNxLIcsM0TNd-O7KS-g2ZsvCuanuJ0,18569
+aixtools-0.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.1.6.dist-info/entry_points.txt,sha256=dHoutULEZx7xXSqJrZdViSVjfInJibfLibi2nRXL3SE,56
+aixtools-0.1.6.dist-info/top_level.txt,sha256=ee4eF-0pqu45zCUVml0mWIhnXQgqMQper2-49BBVHLY,40
+aixtools-0.1.6.dist-info/RECORD,,

tests/unit/vault/test_vault.py

@@ -0,0 +1,114 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from hvac.exceptions import InvalidPath
+
+from aixtools.vault.vault import VaultClient, VaultAuthError
+
+
+@pytest.fixture
+def patched_vault_client():
+    with patch("aixtools.vault.vault.hvac.Client") as mock_hvac_client_cls:
+        fake_hvac_client = MagicMock()
+        fake_hvac_client.is_authenticated.return_value = True
+        mock_hvac_client_cls.return_value = fake_hvac_client
+
+        client = VaultClient()
+        return client
+
+
+@pytest.fixture
+def valid_params():
+    return {
+        "vault_mount_point": "secret",
+        "path_prefix": "path",
+        "env": "dev",
+        "user_id": "test-user",
+        "service_name": "test-service",
+        "user_api_key": "test-api-key"
+    }
+
+
+def test_store_user_service_api_key_success(patched_vault_client, valid_params):
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.return_value = {}
+
+    patched_vault_client.store_user_service_api_key(user_id=valid_params["user_id"],
+                                                    service_name=valid_params['service_name'],
+                                                    user_api_key=valid_params['user_api_key'])
+
+    secret_path = f"{valid_params['path_prefix']}/{valid_params['env']}/{valid_params['user_id']}/{valid_params['service_name']}"
+
+    print("expected secret path", secret_path)
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.assert_called_once_with(
+        secret_path,
+        secret={"user-api-key": valid_params["user_api_key"]},
+        mount_point=valid_params["vault_mount_point"]
+    )
+
+
+def test_store_user_service_api_key_invalid_path(patched_vault_client, valid_params):
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.side_effect = Exception("Invalid Path")
+
+    with pytest.raises(VaultAuthError, match="Invalid Path"):
+        patched_vault_client.store_user_service_api_key(user_id=valid_params["user_id"],
+                                                    service_name=valid_params['service_name'],
+                                                    user_api_key=valid_params['user_api_key'])
+        patched_vault_client.client.assert_not_called()
+
+
+def test_read_user_service_api_key_success(patched_vault_client):
+    """Test successful read of user service API key."""
+    mock_response = {
+        "data": {
+            "data": {
+                "user-api-key": "test-api-key"
+            }
+        }
+    }
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.return_value = mock_response
+
+    result = patched_vault_client.read_user_service_api_key(
+        user_id="test-user",
+        service_name="test-service"
+    )
+
+    assert result == "test-api-key"
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service",
+        mount_point='secret',
+        raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_api_key_secret_not_found(patched_vault_client):
+    """Test read_user_service_api_key when the secret path does not exist."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = InvalidPath
+
+    result = patched_vault_client.read_user_service_api_key(
+        user_id="test-user",
+        service_name="test-service"
+    )
+
+    assert result is None
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service",
+        mount_point="secret",
+        raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_api_key_unexpected_error(patched_vault_client):
+    """Test read_user_service_api_key when an unexpected exception occurs."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = Exception("Unexpected error")
+
+    with pytest.raises(VaultAuthError, match="Unexpected error"):
+        patched_vault_client.read_user_service_api_key(
+            user_id="test-user",
+            service_name="test-service"
+        )
+
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service",
+        mount_point="secret",
+        raise_on_deleted_version=True
+    )

aixtools-0.1.5.dist-info/METADATA

@@ -1,357 +0,0 @@
-Metadata-Version: 2.4
-Name: aixtools
-Version: 0.1.5
-Summary: Tools for AI exploration and debugging
-Requires-Python: >=3.11.2
-Description-Content-Type: text/markdown
-Requires-Dist: a2a-sdk>=0.3.1
-Requires-Dist: cachebox>=5.0.1
-Requires-Dist: chainlit>=2.5.5
-Requires-Dist: colorlog>=6.9.0
-Requires-Dist: fasta2a>=0.5.0
-Requires-Dist: fastmcp>=2.10.2
-Requires-Dist: ipykernel>=6.29.5
-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: pandas>=2.2.3
-Requires-Dist: pydantic-ai>=0.4.10
-Requires-Dist: pylint>=3.3.7
-Requires-Dist: rich>=14.0.0
-Requires-Dist: ruff>=0.11.6
-Requires-Dist: streamlit>=1.44.1
-Requires-Dist: watchdog>=6.0.0
-Provides-Extra: test
-Requires-Dist: pyyaml; extra == "test"
-Provides-Extra: feature
-Requires-Dist: logfire; extra == "feature"
-
-# AiXplore-tools
-
-AiXplore-tools (`axitools` for short) is a Python library to support AI agent development and debugging. It serves two main purposes:
-
-1. **Utility Library**: A collection of utility functions and classes for agent projects
-2. **Debugging Tool**: A Streamlit app (`log_view`) for visualizing and debugging agent applications by viewing messages logged by `ObjectLogger`
-
-
-## `.venv`
-
-```
-# Create a new environment
-uv venv .venv
-source .venv/bin/activate
-
-# Add packages
-uv add pydantic-ai
-uv add mcp[cli]
-uv add chainlit
-uv add streamlit
-
-# Add langchain
-uv add langchain-chroma
-uv add langchain-ollama
-uv add langchain-openai
-
-# Other packages
-uv add colorlog
-uv add pandas
-uv add rich
-uv add ruff
-uv add watchdog
-uv add ipykernel
-
-# Install this code as a package
-uv pip install -e .
-```
-
-## Installation
-
-Installing the package from GitHub, you can use `uv add` (prefered)
-```
-uv add https://github.com/your-org/aixtools.git
-```
-
-or using the `pip` commamnd
-```
-uv pip install https://github.com/your-org/aixtools.git
-```
-
-Remember to set up your [`.env` file](#environment-configuration).
-
-### Updating
-
-This package is in active development and changes often.
-
-You'll want to force uv to re-install or update it to the latest commit, even if the version hasn’t changed.
-
-```
-uv add --upgrade https://github.com/your-org/aixtools.git
-```
-
-### Example: Creatng a new project
-
-```
-# Create a new project
-uv init MyNewProject
-cd MyNewProject
-
-# Add a virtual environemnt and activate it
-uv venv .venv
-source .venv/bin/activate
-
-# Add this package
-uv add https://github.com/your-org/aixtools.git
-```
-
-### Environment Configuration
-
-AIXtools requires specific environment variables to be set for proper operation, especially when working with different model providers.
-
-Here is a [template of an `.env`](./.env_template) file you can use as reference.
-
-#### Setting Up Environment Variables
-
-1. Create a `.env` file in your project root based on the template below:
-
-```
-# Model family (azure, openai, or ollama)
-MODEL_FAMILY=azure
-VDB_EMBEDDINGS_MODEL_FAMILY=azure
-
-MODEL_TIMEOUT=120
-
-# Azure OpenAI
-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
-OPENAI_MODEL_NAME=gpt-4.5-preview
-OPENAI_API_KEY=openai_api_key
-OPENAI_VDB_EMBEDDINGS_MODEL_NAME=text-embedding-3-small
-
-# Ollama models
-OLLAMA_MODEL_NAME=llama3.2:3b-instruct-fp16
-OLLAMA_LOCAL_URL=http://localhost:11434/v1
-OLLAMA_VDB_EMBEDDINGS_MODEL_NAME=snowflake-arctic-embed2:latest
-```
-
-2. Configure the variables according to your needs:
-   - Set `MODEL_FAMILY` to your preferred model provider (`azure`, `openai`, or `ollama`)
-   - Provide the appropriate API keys and endpoints for your chosen provider
-   - Adjust model names and timeouts as needed
-
-#### Important Environment Variables
-
-| Variable | Description | Required |
-|----------|-------------|----------|
-| `MODEL_FAMILY` | The model provider to use (azure, openai, ollama) | Yes |
-| `MODEL_TIMEOUT` | Timeout in seconds for model requests | Yes |
-| `AZURE_*` | Azure OpenAI configuration (if using Azure) | If MODEL_FAMILY=azure |
-| `OPENAI_*` | OpenAI configuration (if using OpenAI) | If MODEL_FAMILY=openai |
-| `OLLAMA_*` | Ollama configuration (if using Ollama) | If MODEL_FAMILY=ollama |
-
-## Debugging tools: `log_view` and `ObjectLogger`
-
-The `ObjectLogger` provides a simple logging facility, these logs can then be visualized with `log_view` app.
-
-
-### `ObjectLogger`
-
-The `ObjectLogger` class allows you to log objects to a file for later analysis. This is particularly useful for debugging agent interactions.
-
-`ObjectLogger` simply serializes objects to a pickle file. It first removes unserializeable parts of the objects (e.g. like `async` managers).
-
-#### Basic Usage
-
-```python
-from aixtools.utils.log_objects import ObjectLogger
-
-# Create a logger as a context manager
-with ObjectLogger() as logger:
-    # Log any pickleable object
-    logger.log({"message": "Hello, world!"})
-    
-    # Log agent responses or any other objects
-    logger.log(agent_response)
-```
-
-#### Logging Agent Interactions
-
-The function `run_agent` logs each node from the agent run, roughtly it looks like this:
-
-```python
-async def run_agent(agent: Agent, prompt: str | list[str]):
-    nodes = []
-    async with agent.iter(prompt) as agent_run:
-        with ObjectLogger(debug=debug) as agent_logger:
-            async for node in agent_run:
-                agent_logger.log(node)
-                nodes.append(node)
-            result = agent_run.result
-    return result.data, nodes
-```
-
-#### Saving Multiple Objects
-
-```python
-from aixtools.utils.log_objects import save_objects_to_logfile
-
-# Save a list of objects to a log file
-objects = [obj1, obj2, obj3]
-save_objects_to_logfile(objects)
-```
-
-### log_view App
-
-The `log_view` app allows you to visualize and analyze the objects logged by `ObjectLogger`.
-
-It is a simple Streamlit app that you can run locally and view the log files.
-
-#### Running the App
-
-You can run the app using the command-line script installed with the package:
-
-```bash
-log_view
-```
-
-Or specify a custom log directory:
-
-```bash
-log_view /path/to/logs
-```
-
-#### Features
-
-The log_view app provides several features for analyzing logged objects:
-
-- **Log File Selection**: Choose from available log files
-- **Filtering**: Filter nodes by text, type, attributes, or regex patterns
-- **Visualization**: Expand/collapse nodes to view details
-- **Export**: Export filtered nodes to JSON
-
-#### Example Workflow
-
-1. Log agent interactions using `ObjectLogger`
-2. Run the `log_view` app to analyze the logs
-3. Filter and explore the logged objects
-4. Export interesting findings for further analysis
-
-## Additional Utilities
-
-AIXtools provides several other utilities to support agent development:
-
-### PersistedDict
-
-A dictionary that persists to a file on disk as JSON or pickle:
-
-```python
-from aixtools.utils.persisted_dict import PersistedDict
-
-# Create a persisted dictionary
-data = PersistedDict("data.json")
-
-# Use it like a regular dictionary
-data["key"] = "value"  # Automatically saved to disk
-```
-
-## Agent Utilities
-
-Functions for creating and running agents with different model providers:
-
-```python
-from aixtools.agents.agent import get_agent, run_agent
-
-# Create an agent with the default model (from MODEL_FAMILY)
-agent = get_agent(system_prompt="You are a helpful assistant.")
-
-# Create an agent with a specific model
-from aixtools.agents.agent import get_model_openai
-model = get_model_openai()
-agent = get_agent(system_prompt="You are a helpful assistant.", model=model)
-
-# Run the agent
-result, nodes = await run_agent(agent, "Tell me about AI")
-```
-
-### Batch Processing
-
-Run multiple agent queries simultaneously:
-
-```python
-from aixtools.agents.agent_batch import agent_batch, AgentQueryParams
-
-# Create query parameters
-query_parameters = [
-    AgentQueryParams(prompt="What is the meaning of life"),
-    AgentQueryParams(prompt="Who is the prime minister of Canada")
-]
-
-# Run queries in batches
-async for result in agent_batch(query_parameters):
-    print(result)
-```
-
-## Project Structure
-
-```
-aixtools/
-├── __init__.py           # Package initialization
-├── log_view.py           # Entry point for log_view app
-├── agents/               # Agent-related functionality
-│   ├── agent.py          # Core agent functions
-│   └── agent_batch.py    # Batch processing for agents
-├── log_view/             # Log viewer application
-│   ├── app.py            # Main Streamlit application
-│   ├── display.py        # Node display utilities
-│   ├── export.py         # Export functionality
-│   ├── filters.py        # Node filtering
-│   ├── log_utils.py      # Log file utilities
-│   └── node_utils.py     # Node processing utilities
-└── utils/                # Utility functions and classes
-    ├── config.py         # Configuration utilities
-    ├── log_objects.py    # Object logging functionality
-    ├── persisted_dict.py # Persistent dictionary
-    └── utils.py          # General utilities
-
-## Logging
-
-The logging system is configured using a standard Python `dictConfig` schema. By default, it will look for a `logging.yaml` or `logging.json` file in your project's root directory. You can also specify a custom path using the `LOGGING_CONFIG_PATH` environment variable.
-
-If no configuration file is found, a default colorized console logger will be used.
-
-### Example `logging.yaml`
-
-```yaml
-version: 1
-disable_existing_loggers: false
-formatters:
-  default:
-    format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-  color:
-    "()": "colorlog.ColoredFormatter"
-    "format": "%(log_color)s%(levelname)-8s%(reset)s [%(name)s] %(message)s"
-handlers:
-  console:
-    class: colorlog.StreamHandler
-    formatter: color
-    level: DEBUG
-  file:
-    class: logging.handlers.RotatingFileHandler
-    formatter: default
-    filename: app.log
-    maxBytes: 10485760 # 10MB
-    backupCount: 5
-    level: INFO
-root:
-  handlers: [console, file]
-  level: DEBUG
-loggers:
-  aixtools:
-    level: INFO
-```