agentex-sdk 0.4.19__py3-none-any.whl → 0.4.21__py3-none-any.whl

agentex/lib/cli/debug/debug_config.py

@@ -4,7 +4,6 @@ Debug configuration models for AgentEx CLI debugging.
 
 import socket
 from enum import Enum
-from typing import Optional
 
 from agentex.lib.utils.model_utils import BaseModel
 
@@ -62,7 +61,7 @@ class DebugConfig(BaseModel):
     def create_both_debug(
         cls, 
         worker_port: int = 5678,
-        acp_port: int = 5679,
+        _acp_port: int = 5679,
         wait_for_attach: bool = False,
         auto_port: bool = True
     ) -> "DebugConfig":

agentex/lib/cli/debug/debug_handlers.py

@@ -4,19 +4,21 @@ Debug process handlers for AgentEx CLI.
 Provides debug-enabled versions of ACP server and temporal worker startup.
 """
 
-import asyncio
 import sys
+import asyncio
+import asyncio.subprocess
+from typing import TYPE_CHECKING, Dict
 from pathlib import Path
-from typing import Dict, TYPE_CHECKING
 
 from rich.console import Console
 
 if TYPE_CHECKING:
-    import asyncio.subprocess
+    pass
 
-from .debug_config import DebugConfig, resolve_debug_port
 from agentex.lib.utils.logging import make_logger
 
+from .debug_config import DebugConfig, resolve_debug_port
+
 logger = make_logger(__name__)
 console = Console()
 

agentex/lib/cli/handlers/agent_handlers.py

@@ -2,14 +2,13 @@ from __future__ import annotations
 
 from pathlib import Path
 
-from python_on_whales import DockerException, docker
 from rich.console import Console
+from python_on_whales import DockerException, docker
 
-from agentex.lib.cli.handlers.run_handlers import RunError
-from agentex.lib.cli.handlers.run_handlers import run_agent as _run_agent
 from agentex.lib.cli.debug import DebugConfig
-from agentex.lib.sdk.config.agent_manifest import AgentManifest
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.cli.handlers.run_handlers import RunError, run_agent as _run_agent
+from agentex.lib.sdk.config.agent_manifest import AgentManifest
 
 logger = make_logger(__name__)
 console = Console()
@@ -25,9 +24,9 @@ def build_agent(
     repository_name: str | None,
     platforms: list[str],
     push: bool = False,
-    secret: str = None,
-    tag: str = None,
-    build_args: list[str] = None,
+    secret: str | None = None,
+    tag: str | None = None,
+    build_args: list[str] | None = None,
 ) -> str:
     """Build the agent locally and optionally push to registry
 
@@ -66,14 +65,14 @@ def build_agent(
         # Log build context information for debugging
         logger.info(f"Build context path: {build_context.path}")
         logger.info(
-            f"Dockerfile path: {build_context.path / build_context.dockerfile_path}"
+            f"Dockerfile path: {build_context.path / build_context.dockerfile_path}"  # type: ignore[operator]
         )
 
         try:
             # Prepare build arguments
             docker_build_kwargs = {
                 "context_path": str(build_context.path),
-                "file": str(build_context.path / build_context.dockerfile_path),
+                "file": str(build_context.path / build_context.dockerfile_path),  # type: ignore[operator]
                 "tags": [image_name],
                 "platforms": platforms,
             }
@@ -129,23 +128,23 @@ def build_agent(
 
 def run_agent(manifest_path: str, debug_config: "DebugConfig | None" = None):
     """Run an agent locally from the given manifest"""
-    import asyncio
-    import signal
     import sys
+    import signal
+    import asyncio
 
     # Flag to track if we're shutting down
     shutting_down = False
 
-    def signal_handler(signum, frame):
+    def signal_handler(signum, _frame):
         """Handle signals by raising KeyboardInterrupt"""
         nonlocal shutting_down
         if shutting_down:
             # If we're already shutting down and get another signal, force exit
-            print(f"\nForce exit on signal {signum}")
+            logger.info(f"Force exit on signal {signum}")
             sys.exit(1)
-        
+
         shutting_down = True
-        print(f"\nReceived signal {signum}, shutting down...")
+        logger.info(f"Received signal {signum}, shutting down...")
         raise KeyboardInterrupt()
     
     # Set up signal handling for the main thread
@@ -155,7 +154,7 @@ def run_agent(manifest_path: str, debug_config: "DebugConfig | None" = None):
     try:
         asyncio.run(_run_agent(manifest_path, debug_config))
     except KeyboardInterrupt:
-        print("Shutdown completed.")
+        logger.info("Shutdown completed.")
         sys.exit(0)
     except RunError as e:
         raise RuntimeError(str(e)) from e

agentex/lib/cli/handlers/cleanup_handlers.py

@@ -1,5 +1,5 @@
-import asyncio
 import os
+import asyncio
 
 from rich.console import Console
 

agentex/lib/cli/handlers/deploy_handlers.py

@@ -25,20 +25,14 @@ AGENTEX_AGENTS_HELM_CHART_VERSION = "0.1.7-v4-beta"
 
 
 class InputDeployOverrides(BaseModel):
-    repository: str | None = Field(
-        default=None, description="Override the repository for deployment"
-    )
-    image_tag: str | None = Field(
-        default=None, description="Override the image tag for deployment"
-    )
+    repository: str | None = Field(default=None, description="Override the repository for deployment")
+    image_tag: str | None = Field(default=None, description="Override the image tag for deployment")
 
 
 def check_helm_installed() -> bool:
     """Check if helm is installed and available"""
     try:
-        result = subprocess.run(
-            ["helm", "version", "--short"], capture_output=True, text=True, check=True
-        )
+        result = subprocess.run(["helm", "version", "--short"], capture_output=True, text=True, check=True)
         logger.info(f"Helm version: {result.stdout.strip()}")
         return True
     except (subprocess.CalledProcessError, FileNotFoundError):
@@ -49,9 +43,7 @@ def add_helm_repo(helm_repository_name: str, helm_repository_url: str) -> None:
     """Add the agentex helm repository if not already added"""
     try:
         # Check if repo already exists
-        result = subprocess.run(
-            ["helm", "repo", "list"], capture_output=True, text=True, check=True
-        )
+        result = subprocess.run(["helm", "repo", "list"], capture_output=True, text=True, check=True)
 
         if helm_repository_name not in result.stdout:
             console.print("Adding agentex helm repository...")
@@ -75,7 +67,6 @@ def add_helm_repo(helm_repository_name: str, helm_repository_url: str) -> None:
         raise HelmError(f"Failed to add helm repository: {e}") from e
 
 
-
 def convert_env_vars_dict_to_list(env_vars: dict[str, str]) -> list[dict[str, str]]:
     """Convert a dictionary of environment variables to a list of dictionaries"""
     return [{"name": key, "value": value} for key, value in env_vars.items()]
@@ -169,11 +160,11 @@ def merge_deployment_configs(
     # Priority: manifest -> environments.yaml -> secrets (highest)
     all_env_vars: dict[str, str] = {}
     secret_env_vars: list[dict[str, str]] = []
-    
+
     # Start with agent_config env vars from manifest
     if agent_config.env:
         all_env_vars.update(agent_config.env)
-    
+
     # Override with environment config env vars if they exist
     if agent_env_config and agent_env_config.helm_overrides and "env" in agent_env_config.helm_overrides:
         env_overrides = agent_env_config.helm_overrides["env"]
@@ -185,8 +176,6 @@ def merge_deployment_configs(
                     env_override_dict[str(env_var["name"])] = str(env_var["value"])
             all_env_vars.update(env_override_dict)
 
-
-
     # Handle credentials and check for conflicts
     if agent_config.credentials:
         for credential in agent_config.credentials:
@@ -199,7 +188,7 @@ def merge_deployment_configs(
                 env_var_name = credential.env_var_name
                 secret_name = credential.secret_name
                 secret_key = credential.secret_key
-            
+
             # Check if the environment variable name conflicts with existing env vars
             if env_var_name in all_env_vars:
                 logger.warning(
@@ -208,7 +197,7 @@ def merge_deployment_configs(
                 )
                 # Remove from regular env vars since secret takes precedence
                 del all_env_vars[env_var_name]
-            
+
             secret_env_vars.append(
                 {
                     "name": env_var_name,
@@ -222,13 +211,14 @@ def merge_deployment_configs(
         # Add auth principal env var if environment config is set
         if agent_env_config.auth:
             from agentex.lib.cli.utils.auth_utils import _encode_principal_context_from_env_config
+
             encoded_principal = _encode_principal_context_from_env_config(agent_env_config.auth)
             logger.info(f"Encoding auth principal from {agent_env_config.auth}")
             if encoded_principal:
                 all_env_vars[EnvVarKeys.AUTH_PRINCIPAL_B64.value] = encoded_principal
             else:
                 raise DeploymentError(f"Auth principal unable to be encoded for agent_env_config: {agent_env_config}")
-        
+
         logger.info(f"Defined agent helm overrides: {agent_env_config.helm_overrides}")
         logger.info(f"Before-merge helm values: {helm_values}")
         if agent_env_config.helm_overrides:
@@ -239,7 +229,7 @@ def merge_deployment_configs(
     # Environment variable precedence: manifest -> environments.yaml -> secrets (highest)
     if all_env_vars:
         helm_values["env"] = convert_env_vars_dict_to_list(all_env_vars)
-    
+
     if secret_env_vars:
         helm_values["secretEnvVars"] = secret_env_vars
 
@@ -252,30 +242,25 @@ def merge_deployment_configs(
 
     # Handle image pull secrets
     if manifest.deployment and manifest.deployment.imagePullSecrets:
-        pull_secrets = [
-            pull_secret.model_dump()
-            for pull_secret in manifest.deployment.imagePullSecrets
-        ]
+        pull_secrets = [pull_secret.model_dump() for pull_secret in manifest.deployment.imagePullSecrets]
         helm_values["global"]["imagePullSecrets"] = pull_secrets
         helm_values["imagePullSecrets"] = pull_secrets
 
     # Add dynamic ACP command based on manifest configuration if command is not set in helm overrides
-    helm_overrides_command = agent_env_config and agent_env_config.helm_overrides and "command" in agent_env_config.helm_overrides
+    helm_overrides_command = (
+        agent_env_config and agent_env_config.helm_overrides and "command" in agent_env_config.helm_overrides
+    )
     if not helm_overrides_command:
         add_acp_command_to_helm_values(helm_values, manifest, manifest_path)
-    
-    console.print("Deploying with the following helm values: ", helm_values)
+
+    logger.info("Deploying with the following helm values: %s", helm_values)
     return helm_values
 
 
 def _deep_merge(base_dict: dict[str, Any], override_dict: dict[str, Any]) -> None:
     """Deep merge override_dict into base_dict"""
     for key, value in override_dict.items():
-        if (
-            key in base_dict
-            and isinstance(base_dict[key], dict)
-            and isinstance(value, dict)
-        ):
+        if key in base_dict and isinstance(base_dict[key], dict) and isinstance(value, dict):
             _deep_merge(base_dict[key], value)
         else:
             base_dict[key] = value
@@ -394,12 +379,8 @@ def deploy_agent(
 
         # Show success message with helpful commands
         console.print("\n[green]🎉 Deployment completed successfully![/green]")
-        console.print(
-            f"[blue]Check deployment status:[/blue] helm status {release_name} -n {namespace}"
-        )
-        console.print(
-            f"[blue]View logs:[/blue] kubectl logs -l app.kubernetes.io/name=agentex-agent -n {namespace}"
-        )
+        console.print(f"[blue]Check deployment status:[/blue] helm status {release_name} -n {namespace}")
+        console.print(f"[blue]View logs:[/blue] kubectl logs -l app.kubernetes.io/name=agentex-agent -n {namespace}")
 
     except subprocess.CalledProcessError as e:
         raise HelmError(

agentex/lib/cli/handlers/run_handlers.py

@@ -1,26 +1,21 @@
-import asyncio
 import os
 import sys
+import asyncio
 from pathlib import Path
 
-from rich.console import Console
 from rich.panel import Panel
+from rich.console import Console
 
-from agentex.lib.cli.handlers.cleanup_handlers import (
-    cleanup_agent_workflows,
-    should_cleanup_on_restart
-)
+# Import debug functionality
+from agentex.lib.cli.debug import DebugConfig, start_acp_server_debug, start_temporal_worker_debug
+from agentex.lib.utils.logging import make_logger
 from agentex.lib.cli.utils.path_utils import (
     get_file_paths,
     calculate_uvicorn_target_for_local,
 )
-
 from agentex.lib.environment_variables import EnvVarKeys
 from agentex.lib.sdk.config.agent_manifest import AgentManifest
-
-# Import debug functionality
-from agentex.lib.cli.debug import DebugConfig, start_acp_server_debug, start_temporal_worker_debug
-from agentex.lib.utils.logging import make_logger
+from agentex.lib.cli.handlers.cleanup_handlers import cleanup_agent_workflows, should_cleanup_on_restart
 
 logger = make_logger(__name__)
 console = Console()
@@ -242,6 +237,8 @@ async def start_temporal_worker(
 async def stream_process_output(process: asyncio.subprocess.Process, prefix: str):
     """Stream process output with prefix"""
     try:
+        if process.stdout is None:
+            return
         while True:
             line = await process.stdout.readline()
             if not line:
@@ -297,11 +294,11 @@ async def run_agent(manifest_path: str, debug_config: "DebugConfig | None" = Non
         manifest_dir = Path(manifest_path).parent
         if debug_config and debug_config.should_debug_acp():
             acp_process = await start_acp_server_debug(
-                file_paths["acp"], manifest.local_development.agent.port, agent_env, debug_config
+                file_paths["acp"], manifest.local_development.agent.port, agent_env, debug_config  # type: ignore[union-attr]
             )
         else:
             acp_process = await start_acp_server(
-                file_paths["acp"], manifest.local_development.agent.port, agent_env, manifest_dir
+                file_paths["acp"], manifest.local_development.agent.port, agent_env, manifest_dir  # type: ignore[union-attr]
             )
         process_manager.add_process(acp_process)
 
@@ -325,7 +322,7 @@ async def run_agent(manifest_path: str, debug_config: "DebugConfig | None" = Non
             tasks.append(worker_task)
 
         console.print(
-            f"\n[green]✓ Agent running at: http://localhost:{manifest.local_development.agent.port}[/green]"
+            f"\n[green]✓ Agent running at: http://localhost:{manifest.local_development.agent.port}[/green]"  # type: ignore[union-attr]
         )
         console.print("[dim]Press Ctrl+C to stop[/dim]\n")
 
@@ -369,8 +366,8 @@ def create_agent_environment(manifest: AgentManifest) -> dict[str, str]:
         "REDIS_URL": "redis://localhost:6379",
         "AGENT_NAME": manifest.agent.name,
         "ACP_TYPE": manifest.agent.acp_type,
-        "ACP_URL": f"http://{manifest.local_development.agent.host_address}",
-        "ACP_PORT": str(manifest.local_development.agent.port),
+        "ACP_URL": f"http://{manifest.local_development.agent.host_address}",  # type: ignore[union-attr]
+        "ACP_PORT": str(manifest.local_development.agent.port),  # type: ignore[union-attr]
     }
 
     # Add authorization principal if set - for local development, auth is optional

agentex/lib/cli/handlers/secret_handlers.py

@@ -1,28 +1,19 @@
-import base64
 import json
-from collections import defaultdict
-from pathlib import Path
+import base64
 from typing import Any
+from pathlib import Path
+from collections import defaultdict
 
-import questionary
-import typer
 import yaml
-from kubernetes.client.rest import ApiException
+import typer
+import questionary
 from rich.console import Console
+from kubernetes.client.rest import ApiException
 
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.types.credentials import CredentialMapping
 from agentex.lib.cli.utils.cli_utils import handle_questionary_cancellation
 from agentex.lib.cli.utils.kubectl_utils import get_k8s_client
-from agentex.lib.cli.utils.kubernetes_secrets_utils import (
-    KUBERNETES_SECRET_TO_MANIFEST_KEY,
-    KUBERNETES_SECRET_TYPE_DOCKERCONFIGJSON,
-    KUBERNETES_SECRET_TYPE_OPAQUE,
-    VALID_SECRET_TYPES,
-    create_image_pull_secret_with_data,
-    create_secret_with_data,
-    get_secret_data,
-    update_image_pull_secret_with_data,
-    update_secret_with_data,
-)
 from agentex.lib.sdk.config.agent_config import AgentConfig
 from agentex.lib.sdk.config.agent_manifest import AgentManifest
 from agentex.lib.sdk.config.deployment_config import (
@@ -30,8 +21,17 @@ from agentex.lib.sdk.config.deployment_config import (
     ImagePullSecretConfig,
     InjectedSecretsValues,
 )
-from agentex.lib.types.credentials import CredentialMapping
-from agentex.lib.utils.logging import make_logger
+from agentex.lib.cli.utils.kubernetes_secrets_utils import (
+    VALID_SECRET_TYPES,
+    KUBERNETES_SECRET_TYPE_OPAQUE,
+    KUBERNETES_SECRET_TO_MANIFEST_KEY,
+    KUBERNETES_SECRET_TYPE_DOCKERCONFIGJSON,
+    get_secret_data,
+    create_secret_with_data,
+    update_secret_with_data,
+    create_image_pull_secret_with_data,
+    update_image_pull_secret_with_data,
+)
 
 logger = make_logger(__name__)
 console = Console()
@@ -131,16 +131,16 @@ def interactive_secret_input(secret_name: str, secret_key: str) -> str:
         return handle_questionary_cancellation(result, "text input")
 
 
-def get_secret(name: str, namespace: str, context: str | None = None) -> dict:
+def get_secret(name: str, namespace: str, context: str | None = None) -> dict[str, Any]:
     """Get details about a secret"""
     v1 = get_k8s_client(context)
 
     try:
         secret = v1.read_namespaced_secret(name=name, namespace=namespace)
         return {
-            "name": secret.metadata.name,
+            "name": secret.metadata.name,  # type: ignore[union-attr]
             "namespace": namespace,
-            "created": secret.metadata.creation_timestamp.isoformat(),
+            "created": secret.metadata.creation_timestamp.isoformat(),  # type: ignore[union-attr]
             "exists": True,
         }
     except ApiException as e:
@@ -173,7 +173,7 @@ def delete_secret(name: str, namespace: str, context: str | None = None) -> None
 
 def get_kubernetes_secrets_by_type(
     namespace: str, context: str | None = None
-) -> dict[str, list[dict]]:
+) -> dict[str, list[dict[str, Any]]]:
     """List metadata about secrets in the namespace"""
     v1 = get_k8s_client(context)
 
@@ -218,7 +218,7 @@ def sync_user_defined_secrets(
     cluster_secret_names = {secret["name"] for secret in found_secrets}
     # Get the secrets from the manifest
     agent_config: AgentConfig = manifest_obj.agent
-    manifest_credentials: list[CredentialMapping] = agent_config.credentials or []
+    manifest_credentials: list[CredentialMapping] = agent_config.credentials or []  # type: ignore[assignment]
 
     if not manifest_credentials:
         console.print("[yellow]No credentials found in manifest[/yellow]")
@@ -465,7 +465,7 @@ def sync_image_pull_secrets(
     }
 
     # Get the secrets from the manifest
-    deployment_config: DeploymentConfig = manifest_obj.deployment
+    deployment_config: DeploymentConfig = manifest_obj.deployment  # type: ignore[assignment]
     manifest_image_pull_secrets: list[ImagePullSecretConfig] = (
         deployment_config.imagePullSecrets or []
     )

agentex/lib/cli/templates/temporal/project/acp.py.j2

@@ -5,23 +5,26 @@ import sys
 if os.getenv("AGENTEX_DEBUG_ENABLED") == "true":
     try:
         import debugpy
+        from agentex.lib.utils.logging import make_logger
+
+        logger = make_logger(__name__)
         debug_port = int(os.getenv("AGENTEX_DEBUG_PORT", "5679"))
         debug_type = os.getenv("AGENTEX_DEBUG_TYPE", "acp")
         wait_for_attach = os.getenv("AGENTEX_DEBUG_WAIT_FOR_ATTACH", "false").lower() == "true"
-        
+
         # Configure debugpy
         debugpy.configure(subProcess=False)
         debugpy.listen(debug_port)
-        
-        print(f"🐛 [{debug_type.upper()}] Debug server listening on port {debug_port}")
-        
+
+        logger.info(f"🐛 [{debug_type.upper()}] Debug server listening on port {debug_port}")
+
         if wait_for_attach:
-            print(f"⏳ [{debug_type.upper()}] Waiting for debugger to attach...")
+            logger.info(f"⏳ [{debug_type.upper()}] Waiting for debugger to attach...")
             debugpy.wait_for_client()
-            print(f"✅ [{debug_type.upper()}] Debugger attached!")
+            logger.info(f"✅ [{debug_type.upper()}] Debugger attached!")
         else:
-            print(f"📡 [{debug_type.upper()}] Ready for debugger attachment")
-            
+            logger.info(f"📡 [{debug_type.upper()}] Ready for debugger attachment")
+
     except ImportError:
         print("❌ debugpy not available. Install with: pip install debugpy")
         sys.exit(1)

agentex/lib/cli/utils/auth_utils.py

@@ -1,13 +1,14 @@
-import base64
 import json
+import base64
 from typing import Any, Dict
 
 from agentex.lib.sdk.config.agent_manifest import AgentManifest
 from agentex.lib.sdk.config.environment_config import AgentAuthConfig
 
+
 # DEPRECATED: Old function for backward compatibility
 # Will be removed in future version
-def _encode_principal_context(manifest: AgentManifest) -> str | None:
+def _encode_principal_context(manifest: AgentManifest) -> str | None:  # noqa: ARG001
     """
     DEPRECATED: This function is deprecated as AgentManifest no longer contains auth.
     Use _encode_principal_context_from_env_config instead.

agentex/lib/cli/utils/credential_utils.py

@@ -1,7 +1,7 @@
 import subprocess
 
+from rich.prompt import Prompt, Confirm
 from rich.console import Console
-from rich.prompt import Confirm, Prompt
 
 from agentex.lib.types.credentials import CredentialMapping
 

agentex/lib/cli/utils/kubectl_utils.py

@@ -1,11 +1,11 @@
 import subprocess
 
 from kubernetes import client, config
-from kubernetes.client.rest import ApiException
 from rich.console import Console
+from kubernetes.client.rest import ApiException
 
-from agentex.lib.cli.utils.exceptions import DeploymentError
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.cli.utils.exceptions import DeploymentError
 
 logger = make_logger(__name__)
 console = Console()
@@ -59,7 +59,7 @@ def list_available_contexts() -> list[str]:
     """List all available kubectl contexts"""
     try:
         contexts, _ = config.list_kube_config_contexts()
-        return [ctx["name"] for ctx in contexts]
+        return [ctx["name"] for ctx in contexts]  # type: ignore[index]
     except Exception as e:
         raise DeploymentError(f"Failed to list kubectl contexts: {e}") from e
 

agentex/lib/cli/utils/kubernetes_secrets_utils.py

@@ -1,11 +1,11 @@
 import base64
 
 from kubernetes import client
-from kubernetes.client.rest import ApiException
 from rich.console import Console
+from kubernetes.client.rest import ApiException
 
-from agentex.lib.cli.utils.kubectl_utils import get_k8s_client
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.cli.utils.kubectl_utils import get_k8s_client
 
 logger = make_logger(__name__)
 console = Console()
@@ -172,11 +172,11 @@ def get_secret_data(
     v1 = get_k8s_client(context)
     try:
         secret = v1.read_namespaced_secret(name=name, namespace=namespace)
-        if secret.data:
+        if secret.data:  # type: ignore[union-attr]
             # Decode base64 data
             return {
                 key: base64.b64decode(value).decode("utf-8")
-                for key, value in secret.data.items()
+                for key, value in secret.data.items()  # type: ignore[union-attr]
             }
         return {}
     except ApiException as e:

agentex/lib/cli/utils/path_utils.py

@@ -1,8 +1,8 @@
-from pathlib import Path
 from typing import Dict
+from pathlib import Path
 
-from agentex.lib.sdk.config.agent_manifest import AgentManifest
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.sdk.config.agent_manifest import AgentManifest
 
 logger = make_logger(__name__)
 

agentex/lib/core/adapters/llm/adapter_litellm.py

@@ -1,15 +1,17 @@
-from collections.abc import AsyncGenerator, Generator
+from typing import override
+from collections.abc import Generator, AsyncGenerator
 
 import litellm as llm
 
-from agentex.lib.core.adapters.llm.port import LLMGateway
-from agentex.lib.types.llm_messages import Completion
 from agentex.lib.utils.logging import make_logger
+from agentex.lib.types.llm_messages import Completion
+from agentex.lib.core.adapters.llm.port import LLMGateway
 
 logger = make_logger(__name__)
 
 
 class LiteLLMGateway(LLMGateway):
+    @override
     def completion(self, *args, **kwargs) -> Completion:
         if kwargs.get("stream", True):
             raise ValueError(
@@ -19,6 +21,7 @@ class LiteLLMGateway(LLMGateway):
         response = llm.completion(*args, **kwargs)
         return Completion.model_validate(response)
 
+    @override
     def completion_stream(self, *args, **kwargs) -> Generator[Completion, None, None]:
         if not kwargs.get("stream"):
             raise ValueError("To use streaming, please set stream=True in the kwargs")
@@ -26,6 +29,7 @@ class LiteLLMGateway(LLMGateway):
         for chunk in llm.completion(*args, **kwargs):
             yield Completion.model_validate(chunk)
 
+    @override
     async def acompletion(self, *args, **kwargs) -> Completion:
         if kwargs.get("stream", True):
             raise ValueError(
@@ -36,11 +40,12 @@ class LiteLLMGateway(LLMGateway):
         response = await llm.acompletion(*args, **kwargs)
         return Completion.model_validate(response)
 
+    @override
     async def acompletion_stream(
         self, *args, **kwargs
     ) -> AsyncGenerator[Completion, None]:
         if not kwargs.get("stream"):
             raise ValueError("To use streaming, please set stream=True in the kwargs")
 
-        async for chunk in await llm.acompletion(*args, **kwargs):
+        async for chunk in await llm.acompletion(*args, **kwargs):  # type: ignore[misc]
             yield Completion.model_validate(chunk)