kiwi-code 0.0.4__py3-none-any.whl

kiwi_tui/client.py

@@ -0,0 +1,539 @@
+"""Autobots client wrapper for API interactions."""
+
+from datetime import datetime, timedelta
+from typing import Optional, Callable
+import time
+import json
+import codecs
+import httpx
+from loguru import logger
+from autobots_client import Client, AuthenticatedClient
+from autobots_client.api.hello import hello_v1_hello_get
+from autobots_client.api.auth import (
+    return_user_and_session_v1_auth_post,
+    refresh_password_email_v1_auth_session_refresh_post,
+)
+from autobots_client.api.actions import async_run_action_v1_actions_id_async_run_post
+from autobots_client.api.action_results import get_action_result_v1_action_results_id_get
+from autobots_client.models.body_return_user_and_session_v1_auth_post import (
+    BodyReturnUserAndSessionV1AuthPost,
+)
+from autobots_client.models.body_async_run_action_v1_actions_id_async_run_post import (
+    BodyAsyncRunActionV1ActionsIdAsyncRunPost,
+)
+from autobots_client.models.body_async_run_action_v1_actions_id_async_run_post_input import (
+    BodyAsyncRunActionV1ActionsIdAsyncRunPostInput,
+)
+
+from .models import AuthTokens, LoginCredentials
+
+
+class AutobotsClientWrapper:
+    """Wrapper around autobots_client for API calls."""
+
+    def __init__(
+        self,
+        base_url: str,
+        access_token: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ):
+        """Initialize autobots client wrapper.
+
+        Args:
+            base_url: Base URL of the autobots backend
+            access_token: Optional access token for authentication
+            api_key: Optional API key for authentication (legacy)
+        """
+        self.base_url = base_url
+        self.access_token = access_token
+        self.api_key = api_key
+
+        # Use access_token first, then api_key, then unauthenticated
+        token = access_token or api_key
+
+        if token:
+            self.client = AuthenticatedClient(
+                base_url=base_url,
+                token=token,
+                raise_on_unexpected_status=False,
+            )
+            logger.info(f"Initialized authenticated client for {base_url}")
+        else:
+            self.client = Client(
+                base_url=base_url,
+                raise_on_unexpected_status=False,
+            )
+            logger.info(f"Initialized unauthenticated client for {base_url}")
+
+    def update_token(self, access_token: str) -> None:
+        """Update the client with a new access token.
+
+        Args:
+            access_token: New access token
+        """
+        self.access_token = access_token
+        self.client = AuthenticatedClient(
+            base_url=self.base_url,
+            token=access_token,
+            raise_on_unexpected_status=False,
+        )
+        logger.info("Client updated with new access token")
+
+    def hello(self) -> tuple[bool, str]:
+        """Call the hello API endpoint.
+
+        Returns:
+            Tuple of (success, message)
+        """
+        try:
+            logger.debug("Calling hello API")
+            response = hello_v1_hello_get.sync_detailed(client=self.client)
+
+            if response.status_code == 200:
+                # Get response content as text
+                message = response.content.decode('utf-8') if response.content else "Hello from Autobots!"
+                logger.info(f"Hello API success: {message}")
+                return True, message
+            else:
+                error_msg = f"Hello API returned status {response.status_code}"
+                logger.warning(error_msg)
+                return False, error_msg
+
+        except Exception as e:
+            error_msg = f"Failed to call hello API: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg
+
+    def login(self, credentials: LoginCredentials) -> tuple[bool, Optional[AuthTokens], str]:
+        """Login with username and password.
+
+        Args:
+            credentials: Login credentials
+
+        Returns:
+            Tuple of (success, auth_tokens, message)
+        """
+        try:
+            logger.info(f"Attempting login for user: {credentials.username}")
+
+            body = BodyReturnUserAndSessionV1AuthPost(
+                username=credentials.username,
+                password=credentials.password,
+            )
+
+            response = return_user_and_session_v1_auth_post.sync_detailed(
+                client=self.client,
+                body=body,
+            )
+
+            if response.status_code == 200 and response.parsed:
+                auth_response = response.parsed
+
+                if auth_response.session:
+                    session = auth_response.session
+
+                    # Calculate expiry time
+                    expires_at = None
+                    if hasattr(session, 'expires_in') and session.expires_in:
+                        expires_at = datetime.now() + timedelta(seconds=session.expires_in)
+
+                    tokens = AuthTokens(
+                        access_token=session.access_token,
+                        refresh_token=session.refresh_token,
+                        token_type=session.token_type,
+                        expires_at=expires_at,
+                    )
+
+                    # Update client with new token
+                    self.update_token(tokens.access_token)
+
+                    logger.info(f"Login successful for user: {credentials.username}")
+                    return True, tokens, "Login successful"
+                else:
+                    error_msg = "No session in response"
+                    logger.error(error_msg)
+                    return False, None, error_msg
+            else:
+                error_msg = f"Login failed with status {response.status_code}"
+                if response.content:
+                    error_msg += f": {response.content.decode('utf-8')}"
+                logger.error(error_msg)
+                return False, None, error_msg
+
+        except Exception as e:
+            error_msg = f"Login error: {str(e)}"
+            logger.error(error_msg)
+            return False, None, error_msg
+
+    def refresh_token(self, refresh_token: str) -> tuple[bool, Optional[AuthTokens], str]:
+        """Refresh access token using refresh token.
+
+        Args:
+            refresh_token: Refresh token
+
+        Returns:
+            Tuple of (success, new_auth_tokens, message)
+        """
+        try:
+            logger.info("Attempting to refresh access token")
+
+            # Need authenticated client for refresh
+            if not isinstance(self.client, AuthenticatedClient):
+                temp_client = AuthenticatedClient(
+                    base_url=self.base_url,
+                    token=refresh_token,
+                    raise_on_unexpected_status=False,
+                )
+            else:
+                temp_client = self.client
+
+            response = refresh_password_email_v1_auth_session_refresh_post.sync_detailed(
+                client=temp_client,
+                refresh_token=refresh_token,
+            )
+
+            if response.status_code == 200 and response.parsed:
+                auth_response = response.parsed
+
+                if auth_response.session:
+                    session = auth_response.session
+
+                    # Calculate expiry time
+                    expires_at = None
+                    if hasattr(session, 'expires_in') and session.expires_in:
+                        expires_at = datetime.now() + timedelta(seconds=session.expires_in)
+
+                    tokens = AuthTokens(
+                        access_token=session.access_token,
+                        refresh_token=session.refresh_token,
+                        token_type=session.token_type,
+                        expires_at=expires_at,
+                    )
+
+                    # Update client with new token
+                    self.update_token(tokens.access_token)
+
+                    logger.info("Token refresh successful")
+                    return True, tokens, "Token refreshed successfully"
+                else:
+                    error_msg = "No session in refresh response"
+                    logger.error(error_msg)
+                    return False, None, error_msg
+            else:
+                error_msg = f"Token refresh failed with status {response.status_code}"
+                logger.error(error_msg)
+                return False, None, error_msg
+
+        except Exception as e:
+            error_msg = f"Token refresh error: {str(e)}"
+            logger.error(error_msg)
+            return False, None, error_msg
+
+    def run_action_async(self, action_id: str, user_input: str, action_result_id: Optional[str] = None) -> tuple[bool, Optional[str], str]:
+        """Run an action asynchronously and return the run ID.
+
+        Args:
+            action_id: ID of the action to run
+            user_input: User input text for the action
+            action_result_id: Optional ID to continue an existing conversation
+
+        Returns:
+            Tuple of (success, run_id, message)
+        """
+        try:
+            if action_result_id:
+                logger.info(f"Running action {action_id} with input: {user_input}, continuing run_id: {action_result_id}")
+            else:
+                logger.info(f"Running action {action_id} with input: {user_input}")
+
+            if not isinstance(self.client, AuthenticatedClient):
+                return False, None, "Authentication required"
+
+            # Create input body - DataBlock structure
+            input_data = BodyAsyncRunActionV1ActionsIdAsyncRunPostInput()
+            input_data.additional_properties = {
+                "text": user_input,
+                "urls": [],
+                "metadata": {
+                    "process_url_in_text": False,
+                    "process_attachment_url": True,
+                    "attachment_requested": 0
+                }
+            }
+
+            body = BodyAsyncRunActionV1ActionsIdAsyncRunPost(input_=input_data)
+
+            # Import UNSET for optional parameter
+            from autobots_client.types import UNSET
+
+            response = async_run_action_v1_actions_id_async_run_post.sync_detailed(
+                id=action_id,
+                client=self.client,
+                body=body,
+                action_result_id=action_result_id if action_result_id else UNSET,
+            )
+
+            if response.status_code == 200 and response.parsed:
+                action_result = response.parsed
+
+                # ActionResultDoc has field_id as the run ID
+                run_id = getattr(action_result, 'field_id', None)
+
+                # Fallback to additional_properties if needed
+                if not run_id and hasattr(action_result, 'additional_properties'):
+                    run_id = action_result.additional_properties.get('_id')
+
+                if run_id:
+                    logger.info(f"Action started with run ID: {run_id}")
+                    return True, run_id, f"Action started with ID: {run_id}"
+                else:
+                    error_msg = "No run ID found in response"
+                    logger.error(error_msg)
+                    return False, None, error_msg
+            else:
+                error_msg = f"Failed to start action: status {response.status_code}"
+                logger.error(error_msg)
+                return False, None, error_msg
+
+        except Exception as e:
+            error_msg = f"Error running action: {str(e)}"
+            logger.error(error_msg)
+            return False, None, error_msg
+
+    def get_action_result(self, run_id: str) -> tuple[bool, Optional[dict], str]:
+        """Get the result of an action run.
+
+        Args:
+            run_id: ID of the action run (action_results id)
+
+        Returns:
+            Tuple of (success, result_data, message)
+        """
+        try:
+            if not isinstance(self.client, AuthenticatedClient):
+                return False, None, "Authentication required"
+
+            response = get_action_result_v1_action_results_id_get.sync_detailed(
+                id=run_id,
+                client=self.client,
+            )
+
+            if response.status_code == 200 and response.parsed:
+                result = response.parsed
+
+                # Convert the parsed result to dict preserving full structure
+                # The result is ActionResultDoc which has:
+                # - status: EventResultStatus
+                # - result: ActionDoc (which contains results array)
+                # - error_message: DataBlock
+
+                result_dict = {}
+
+                # Extract status
+                if hasattr(result, 'status'):
+                    status_obj = result.status
+                    if hasattr(status_obj, 'value'):
+                        result_dict["status"] = status_obj.value
+                    else:
+                        result_dict["status"] = str(status_obj)
+
+                # Extract the full ActionDoc with results array
+                if hasattr(result, 'result'):
+                    action_doc = result.result
+                    # Convert ActionDoc to dict
+                    if hasattr(action_doc, 'to_dict'):
+                        result_dict["result"] = action_doc.to_dict()
+                    elif hasattr(action_doc, 'additional_properties'):
+                        result_dict["result"] = action_doc.additional_properties
+                    else:
+                        # Try to extract key fields
+                        action_dict = {}
+                        if hasattr(action_doc, 'results'):
+                            # This is the results array we need
+                            results = action_doc.results
+                            if isinstance(results, list):
+                                action_dict["results"] = [
+                                    r.to_dict() if hasattr(r, 'to_dict') else r
+                                    for r in results
+                                ]
+                            else:
+                                action_dict["results"] = results
+                        result_dict["result"] = action_dict
+
+                # Extract error message
+                if hasattr(result, 'error_message'):
+                    error_obj = result.error_message
+                    if error_obj:
+                        if hasattr(error_obj, 'to_dict'):
+                            result_dict["error"] = error_obj.to_dict()
+                        elif hasattr(error_obj, 'text'):
+                            result_dict["error"] = error_obj.text
+                        else:
+                            result_dict["error"] = str(error_obj)
+
+                result_dict["id"] = run_id
+
+                logger.info(f"Parsed result: status={result_dict.get('status')}, has_result={('result' in result_dict)}")
+                logger.info(f"Result structure keys: {result_dict.keys()}")
+                if "result" in result_dict and isinstance(result_dict["result"], dict):
+                    logger.info(f"Result.result keys: {result_dict['result'].keys()}")
+
+                return True, result_dict, "Result retrieved"
+            else:
+                error_msg = f"Failed to get result: status {response.status_code}"
+                logger.error(error_msg)
+                return False, None, error_msg
+
+        except Exception as e:
+            error_msg = f"Error getting result: {str(e)}"
+            logger.error(error_msg)
+            return False, None, error_msg
+
+    def poll_action_result(self, run_id: str, max_attempts: int = 60, interval: float = 2.0) -> tuple[bool, Optional[dict], str]:
+        """Poll for action result until completion or timeout.
+
+        Args:
+            run_id: ID of the action run
+            max_attempts: Maximum number of polling attempts (default 60 = 2 minutes)
+            interval: Seconds between polls
+
+        Returns:
+            Tuple of (success, result_data, message)
+        """
+        logger.info(f"Polling for result of run {run_id}")
+
+        for attempt in range(max_attempts):
+            success, result, message = self.get_action_result(run_id)
+
+            if not success:
+                return False, None, message
+
+            status = result.get("status", "").lower() if result else ""
+
+            # Check for completion - status can be "completed", "success", or "finished"
+            if status in ["completed", "success", "finished"]:
+                logger.info(f"Action completed after {attempt + 1} attempts")
+                return True, result, "Action completed successfully"
+
+            # Check for failure
+            elif status in ["failed", "error"]:
+                error = result.get("error", "Unknown error") if result else "Unknown error"
+                logger.error(f"Action failed: {error}")
+                return False, result, f"Action failed: {error}"
+
+            # Still processing
+            elif status in ["processing", "running", "pending"]:
+                logger.debug(f"Action still {status}, attempt {attempt + 1}/{max_attempts}")
+                if attempt < max_attempts - 1:
+                    time.sleep(interval)
+
+            # Unknown status
+            else:
+                logger.warning(f"Unknown status '{status}', continuing to poll...")
+                if attempt < max_attempts - 1:
+                    time.sleep(interval)
+
+        return False, result, f"Polling timeout after {max_attempts * interval}s - action may still be running"
+
+    async def stream_action_result(self, run_id: str, on_message: Callable[[dict], None]) -> None:
+        """Stream action result updates via Server-Sent Events.
+
+        Args:
+            run_id: ID of the action run
+            on_message: Callback function to handle incoming messages
+        """
+        sse_endpoint = f"{self.base_url}/v1/server_sent_events/stream/{run_id}"
+        logger.info(f"Connecting to SSE: {sse_endpoint}")
+
+        try:
+            # Create headers
+            headers = {
+                "Accept": "text/event-stream",
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+            }
+            if self.access_token:
+                headers["Authorization"] = f"Bearer {self.access_token}"
+
+            # Use httpx to stream SSE with longer timeout
+            async with httpx.AsyncClient(timeout=None) as client:
+                async with client.stream("GET", sse_endpoint, headers=headers) as response:
+                    logger.info(f"SSE connected, status: {response.status_code}")
+
+                    if response.status_code != 200:
+                        logger.error(f"SSE connection failed with status {response.status_code}")
+                        return
+
+                    # Use incremental UTF-8 decoder to handle multi-byte characters across chunk boundaries
+                    decoder = codecs.getincrementaldecoder("utf-8")()
+                    buffer = ""
+
+                    # Process SSE stream byte by byte with incremental decoder
+                    async for chunk in response.aiter_bytes():
+                        # Decode chunk with incremental decoder (handles partial multi-byte chars)
+                        text_chunk = decoder.decode(chunk, final=False)
+                        buffer += text_chunk
+
+                        # Process complete lines from buffer
+                        while "\n" in buffer:
+                            line, buffer = buffer.split("\n", 1)
+                            line = line.strip()
+
+                            # Skip empty lines
+                            if not line:
+                                continue
+
+                            # Skip keep-alive comments (`:keep-alive`)
+                            if line.startswith(":"):
+                                logger.debug("Received SSE keep-alive")
+                                continue
+
+                            # Parse SSE data format: "data: {...}" or "data: plain text"
+                            if line.startswith("data: "):
+                                data_str = line[6:].strip()
+
+                                # Try to parse as JSON first
+                                try:
+                                    data = json.loads(data_str)
+                                    logger.info(f"SSE JSON received, keys: {data.keys() if isinstance(data, dict) else type(data)}")
+
+                                    # Call the callback with parsed data
+                                    on_message(data)
+
+                                    # Check if action is complete
+                                    if isinstance(data, dict):
+                                        status = data.get("status", "").lower()
+                                        if status in ["completed", "success", "finished", "failed", "error"]:
+                                            logger.info(f"Action completed with status: {status}")
+                                            return
+
+                                except json.JSONDecodeError:
+                                    # Not JSON, treat as plain text status message
+                                    logger.info(f"SSE text message: {data_str}")
+
+                                    # Mark as status so dashboard can distinguish from real output
+                                    text_data = {
+                                        "type": "status",
+                                        "text": data_str,
+                                    }
+                                    on_message(text_data)
+
+                                except Exception as e:
+                                    logger.error(f"Error processing SSE data: {e}")
+
+                    # Flush any remaining bytes in decoder
+                    remaining = decoder.decode(b"", final=True)
+                    if remaining:
+                        buffer += remaining
+
+        except Exception as e:
+            logger.error(f"SSE error: {e}")
+            raise
+
+    def close(self) -> None:
+        """Close the client connection."""
+        try:
+            self.client.__exit__(None, None, None)
+            logger.info("Client connection closed")
+        except Exception as e:
+            logger.warning(f"Error closing client: {e}")