aixtools 0.2.11__tar.gz → 0.2.13__tar.gz

{aixtools-0.2.11 → aixtools-0.2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.11
+Version: 0.2.13
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.2.11 → aixtools-0.2.13}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.11'
-__version_tuple__ = version_tuple = (0, 2, 11)
+__version__ = version = '0.2.13'
+__version_tuple__ = version_tuple = (0, 2, 13)
 
-__commit_id__ = commit_id = 'g66be32dcf'
+__commit_id__ = commit_id = 'gb809b4cc2'

{aixtools-0.2.11 → aixtools-0.2.13}/aixtools/agents/agent.py

@@ -51,17 +51,17 @@ def _get_model_bedrock(model_name=BEDROCK_MODEL_NAME, aws_region=AWS_REGION):
     return BedrockConverseModel(model_name=model_name, provider=provider)
 
 
-def _get_model_ollama(model_name=OLLAMA_MODEL_NAME, ollama_url=OLLAMA_URL):
+def _get_model_ollama(model_name=OLLAMA_MODEL_NAME, ollama_url=OLLAMA_URL, http_client=None):
     assert ollama_url, "OLLAMA_URL is not set"
     assert model_name, "Model name is not set"
-    provider = OpenAIProvider(base_url=ollama_url)
+    provider = OpenAIProvider(base_url=ollama_url, http_client=http_client)
     return OpenAIChatModel(model_name=model_name, provider=provider)
 
 
-def _get_model_openai(model_name=OPENAI_MODEL_NAME, openai_api_key=OPENAI_API_KEY):
+def _get_model_openai(model_name=OPENAI_MODEL_NAME, openai_api_key=OPENAI_API_KEY, http_client=None):
     assert openai_api_key, "OPENAI_API_KEY is not set"
     assert model_name, "Model name is not set"
-    provider = OpenAIProvider(api_key=openai_api_key)
+    provider = OpenAIProvider(api_key=openai_api_key, http_client=http_client)
     return OpenAIChatModel(model_name=model_name, provider=provider)
 
 
@@ -76,22 +76,27 @@ def _get_model_openai_azure(
     assert azure_openai_api_version, "AZURE_OPENAI_API_VERSION is not set"
     assert model_name, "Model name is not set"
     client = AsyncAzureOpenAI(
-        azure_endpoint=azure_openai_endpoint, api_version=azure_openai_api_version, api_key=azure_openai_api_key
+        azure_endpoint=azure_openai_endpoint,
+        api_version=azure_openai_api_version,
+        api_key=azure_openai_api_key,
     )
     return OpenAIChatModel(model_name=model_name, provider=OpenAIProvider(openai_client=client))
 
 
 def _get_model_open_router(
-    model_name=OPENROUTER_MODEL_NAME, openrouter_api_url=OPENROUTER_API_URL, openrouter_api_key=OPENROUTER_API_KEY
+    model_name=OPENROUTER_MODEL_NAME,
+    openrouter_api_url=OPENROUTER_API_URL,
+    openrouter_api_key=OPENROUTER_API_KEY,
+    http_client=None,
 ):
     assert openrouter_api_url, "OPENROUTER_API_URL is not set"
     assert openrouter_api_key, "OPENROUTER_API_KEY is not set"
     assert model_name, "Model name is not set, missing 'OPENROUTER_MODEL_NAME' environment variable?"
-    provider = OpenAIProvider(base_url=openrouter_api_url, api_key=openrouter_api_key)
+    provider = OpenAIProvider(base_url=openrouter_api_url, api_key=openrouter_api_key, http_client=http_client)
     return OpenAIChatModel(model_name, provider=provider)
 
 
-def get_model(model_family=MODEL_FAMILY, model_name=None, **kwargs):
+def get_model(model_family=MODEL_FAMILY, model_name=None, http_client=None, **kwargs):
     """Create and return appropriate model instance based on specified family and name."""
     assert model_family is not None and model_family != "", f"Model family '{model_family}' is not set"
     match model_family:
@@ -100,11 +105,13 @@ def get_model(model_family=MODEL_FAMILY, model_name=None, **kwargs):
         case "bedrock":
             return _get_model_bedrock(model_name=model_name or BEDROCK_MODEL_NAME, **kwargs)
         case "ollama":
-            return _get_model_ollama(model_name=model_name or OLLAMA_MODEL_NAME, **kwargs)
+            return _get_model_ollama(model_name=model_name or OLLAMA_MODEL_NAME, http_client=http_client, **kwargs)
         case "openai":
-            return _get_model_openai(model_name=model_name or OPENAI_MODEL_NAME, **kwargs)
+            return _get_model_openai(model_name=model_name or OPENAI_MODEL_NAME, http_client=http_client, **kwargs)
         case "openrouter":
-            return _get_model_open_router(model_name=model_name or OPENROUTER_MODEL_NAME, **kwargs)
+            return _get_model_open_router(
+                model_name=model_name or OPENROUTER_MODEL_NAME, http_client=http_client, **kwargs
+            )
         case _:
             raise ValueError(f"Model family '{model_family}' not supported")
 
@@ -119,12 +126,15 @@ def get_agent(  # noqa: PLR0913, pylint: disable=too-many-arguments,too-many-pos
     model_settings=None,
     output_type: Any = str,
     deps_type=NoneType,
+    http_client=None,
 ) -> Agent:
     """Get a PydanticAI agent"""
     if model_settings is None:
         model_settings = ModelSettings(timeout=MODEL_TIMEOUT)
     if model is None:
-        model = get_model()
+        model = get_model(
+            http_client=http_client,
+        )
     agent = Agent(
         model=model,
         output_type=output_type,

{aixtools-0.2.11 → aixtools-0.2.13}/aixtools/auth/auth.py

@@ -16,6 +16,7 @@ from mcp.server.auth.provider import (
 from aixtools.utils import config
 
 logger = logging.getLogger(__name__)
+TEST_CLIENT = "test-client"
 
 
 class AuthTokenErrorCode(str, enum.Enum):
@@ -163,8 +164,15 @@ class AccessTokenAuthProvider(AuthProvider):
 
     async def verify_token(self, token: str) -> AccessToken:
         """Verify the access token and return an AccessToken object."""
-
         logger.info("Received verify token request")
+        test_token = config.AUTH_TEST_TOKEN
+
+        # check if the token is a test token
+        # this is used for integration test run
+        if test_token and token == test_token:
+            logger.info("Using test token:%s", test_token)
+            return AccessToken(token=token, client_id=TEST_CLIENT, scopes=[], expires_at=None)
+
         claims = self.token_verifier.verify(token)
         scopes = claims.get("scp", "")
         self.token_verifier.authorize_claims(claims, self.app_scope)

{aixtools-0.2.11 → aixtools-0.2.13}/aixtools/utils/config.py

@@ -124,7 +124,7 @@ 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", allow_empty=True)
-VAULT_ENV = get_variable_env("ENV", allow_empty=True)
+VAULT_ENV = get_variable_env("VAULT_ENV", allow_empty=True)
 VAULT_MOUNT_POINT = get_variable_env("VAULT_MOUNT_POINT", allow_empty=True)
 VAULT_PATH_PREFIX = get_variable_env("VAULT_PATH_PREFIX", allow_empty=True)
 
@@ -142,3 +142,5 @@ APP_AUTHORIZED_GROUPS = get_variable_env("APP_AUTHORIZED_GROUPS", allow_empty=Tr
 # used to skip authorization in local tests if required.
 SKIP_MCP_AUTHORIZATION = str2bool(get_variable_env("SKIP_MCP_AUTHORIZATION", True, False))
 APP_DEFAULT_SCOPE = get_variable_env("APP_DEFAULT_SCOPE", allow_empty=True)
+
+AUTH_TEST_TOKEN = get_variable_env("AUTH_TEST_TOKEN", allow_empty=True)

{aixtools-0.2.11 → aixtools-0.2.13}/pyproject.toml

@@ -71,3 +71,6 @@ ignore = ["_version.py"]
 [tool.setuptools_scm]
 write_to = "aixtools/_version.py"
 version_scheme = "release-branch-semver"
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"