ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/handlers/stream.py

@@ -0,0 +1,337 @@
+from typing import AsyncGenerator, Dict, Any
+import asyncio
+from aiohttp import web
+from datamodel.parsers.json import json_encoder, json_decoder  # pylint: disable=E0611 # noqa
+from navigator.views import BaseHandler
+from navigator_auth.conf import exclude_list
+from parrot.bots import AbstractBot
+
+
+class StreamHandler(BaseHandler):
+    """Streaming Endpoints for Parrot LLM Responses.
+
+    Supports:
+    - SSE (Server-Sent Events)
+    - WebSockets
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.active_connections = set()
+
+    def _get_botmanager(self, request: web.Request):
+        """Retrieve the bot manager from the application context."""
+        try:
+            return request.app['bot_manager']
+        except KeyError as e:
+            raise web.HTTPInternalServerError(
+                reason="Bot manager not found in application."
+            ) from e
+
+    async def _get_bot(self, request: web.Request) -> AbstractBot:
+        """Retrieve the bot instance based on bot_id from the request."""
+        bot_manager = self._get_botmanager(request)
+        bot_id = request.match_info.get('bot_id')
+        bot = await bot_manager.get_bot(bot_id)
+        if bot is None:
+            raise web.HTTPNotFound(
+                reason=f"Bot with ID '{bot_id}' not found."
+            )
+        return bot
+
+    def _extract_stream_params(self, payload: Dict[str, Any], *extra_ignored_keys: str):
+        """Split incoming payload into prompt and kwargs for ask_stream."""
+        ignored_keys = {"prompt", *extra_ignored_keys}
+        prompt = payload.get('prompt', '')
+        kwargs = {k: v for k, v in payload.items() if k not in ignored_keys}
+        return prompt, kwargs
+
+    async def stream_sse(self, request: web.Request) -> web.StreamResponse:
+        """
+        Server-Sent Events (SSE) streaming endpoint
+        Best for: Unidirectional streaming, HTTP/1.1 compatible
+        """
+        data = await request.json()
+        prompt, ask_kwargs = self._extract_stream_params(data)
+        bot = await self._get_bot(request)
+        response = web.StreamResponse(
+            status=200,
+            reason='OK',
+            headers={
+                'Content-Type': 'text/event-stream',
+                'Cache-Control': 'no-cache',
+                'Connection': 'keep-alive',
+                'Access-Control-Allow-Origin': '*',
+                'X-Accel-Buffering': 'no',  # Disable nginx buffering
+            }
+        )
+        await response.prepare(request)
+        try:
+            async for chunk in bot.ask_stream(prompt, **ask_kwargs):
+                # Format as SSE
+                sse_data = f"data: {json_encoder({'content': chunk})}\n\n"
+                await response.write(sse_data.encode('utf-8'))
+                await response.drain()
+
+            # Send completion event
+            await response.write(b"data: [DONE]\n\n")
+            await response.drain()
+        except asyncio.CancelledError as e:
+            raise web.HTTPInternalServerError(
+                reason="Client disconnected during streaming."
+            ) from e
+        except Exception as e:
+            await response.write(
+                f"error: {str(e)}\n\n".encode('utf-8')
+            )
+        finally:
+            await response.write_eof()
+        return response
+
+    async def stream_ndjson(self, request: web.Request) -> web.StreamResponse:
+        """
+        NDJSON (Newline Delimited JSON) streaming endpoint
+        Best for: Clients that can parse JSON lines, more flexible than SSE
+        """
+        data = await request.json()
+        prompt, ask_kwargs = self._extract_stream_params(data)
+        bot = await self._get_bot(request)
+        response = web.StreamResponse(
+            status=200,
+            reason='OK',
+            headers={
+                'Content-Type': 'application/x-ndjson',
+                'Cache-Control': 'no-cache',
+                'Connection': 'keep-alive',
+                'Access-Control-Allow-Origin': '*',
+            }
+        )
+        await response.prepare(request)
+        try:
+            async for chunk in bot.ask_stream(prompt, **ask_kwargs):
+                # Each line is a complete JSON object
+                line = json_encoder({
+                    'type': 'content',
+                    'data': chunk,
+                    'timestamp': asyncio.get_event_loop().time()
+                }) + '\n'
+                await response.write(line.encode('utf-8'))
+                await response.drain()
+
+            # Send completion line
+            await response.write(
+                json_encoder({'done': True}).encode('utf-8') + b'\n'
+            )
+            await response.drain()
+        except asyncio.CancelledError as e:
+            raise web.HTTPInternalServerError(
+                reason="Client disconnected during streaming."
+            ) from e
+        except Exception as e:
+            error_line = json_encoder({'error': str(e)}) + '\n'
+            await response.write(error_line.encode('utf-8'))
+        finally:
+            await response.write_eof()
+        return response
+
+    async def stream_chunked(self, request: web.Request) -> web.StreamResponse:
+        """
+        Plain chunked transfer encoding
+        Best for: Simple text streaming without special formatting
+        """
+        data = await request.json()
+        prompt, ask_kwargs = self._extract_stream_params(data)
+        bot = await self._get_bot(request)
+        response = web.StreamResponse(
+            status=200,
+            reason='OK',
+            headers={
+                'Content-Type': 'text/plain; charset=utf-8',
+                'Transfer-Encoding': 'chunked',
+                'Cache-Control': 'no-cache',
+                'Connection': 'keep-alive',
+                'Access-Control-Allow-Origin': '*',
+            }
+        )
+        await response.prepare(request)
+        try:
+            async for chunk in bot.ask_stream(prompt, **ask_kwargs):
+                await response.write(chunk.encode('utf-8'))
+                await response.drain()
+
+            # Indicate end of stream
+            await response.write_eof()
+        except asyncio.CancelledError as e:
+            raise web.HTTPInternalServerError(
+                reason="Client disconnected during streaming."
+            ) from e
+        except Exception as e:
+            await response.write(f"\n[ERROR]: {str(e)}\n".encode('utf-8'))
+        finally:
+            await response.write_eof()
+        return response
+
+    async def stream_websocket(self, request: web.Request) -> web.WebSocketResponse:
+        """
+        WebSocket endpoint for bidirectional streaming
+        Best for: Real-time bidirectional communication, chat applications
+        """
+        ws = web.WebSocketResponse(
+            heartbeat=30.0,  # Send ping every 30s
+            max_msg_size=10 * 1024 * 1024  # 10MB max message size
+        )
+        # Extract and validate JWT from Sec-WebSocket-Protocol
+        # Client sends: new WebSocket(url, ["jwt", token])
+        # Header received: Sec-WebSocket-Protocol: jwt, <token>
+        protocol_header = request.headers.get('Sec-WebSocket-Protocol')
+        selected_protocol = None
+
+        if protocol_header:
+            parts = [p.strip() for p in protocol_header.split(',')]
+            if 'jwt' in parts:
+                # Find the token (assuming it's the other part)
+                # This is a bit naive if there are other protocols, but fits the requirement
+                try:
+                    token_idx = parts.index('jwt') + 1
+                    # If 'jwt' is the last item, try look elsewhere or it might be (token, jwt) order?
+                    # Browsers usually send in order requested.
+                    # Actually, if we send ["jwt", "token"], header is "jwt, token"
+                    # But we need to identify WHICH one is the token.
+                    # Heuristic: The token is the long string that isn't 'jwt'.
+                    # Or simpler: remove 'jwt' and take the first remaining non-empty string.
+                    parts.remove('jwt')
+                    if parts:
+                        token = parts[0]
+                        if await self._validate_token(token):
+                            # Auth success
+                            # We MUST return the selected protocol in the response
+                            # to the client, otherwise the WS connection fails.
+                            # Usually we echo 'jwt' or the protocol used.
+                            selected_protocol = 'jwt'
+                        else:
+                            raise web.HTTPUnauthorized(reason="Invalid Token")
+                    else:
+                        raise web.HTTPUnauthorized(reason="Missing Token")
+                except (ValueError, IndexError):
+                     raise web.HTTPUnauthorized(reason="Invalid Protocol Format")
+
+        if selected_protocol:
+             await ws.prepare(request, protocols=[selected_protocol])
+        else:
+             await ws.prepare(request)
+
+        self.active_connections.add(ws)
+        bot = await self._get_bot(request)
+
+        try:
+            await ws.send_json({
+                'type': 'connection',
+                'status': 'connected',
+                'message': 'WebSocket connection established'
+            })
+
+            async for msg in ws:
+                if msg.type == web.WSMsgType.TEXT:
+                    # Handle incoming messages
+                    try:
+                        data = json_decoder(msg.data)
+                        await self._handle_message(ws, data, bot)
+                    except Exception:
+                        await ws.send_json({
+                            'type': 'error',
+                            'message': 'Invalid JSON'
+                        })
+                elif msg.type == web.WSMsgType.ERROR:
+                    # Handle errors
+                    print(f'WebSocket error: {ws.exception()}')
+                    self.active_connections.remove(ws)
+        except Exception as e:
+            self.active_connections.remove(ws)
+            raise web.HTTPInternalServerError(
+                reason="Error occurred during WebSocket communication."
+            ) from e
+        return ws
+
+    async def _validate_token(self, token: str) -> bool:
+        """Validate the provided token using the app's auth system."""
+        if not token:
+            return False
+        # auth_manager = self.app.get('auth_manager')
+        # if not auth_manager:
+        #     return False
+        # is_valid = await auth_manager.validate_token(token)
+        # return is_valid
+        return True  # Temporarily allow all tokens for testing
+
+    async def _handle_message(self, ws: web.WebSocketResponse, data: dict, bot: AbstractBot):
+        """Handle incoming WebSocket messages"""
+        msg_type = data.get('type')
+        if msg_type == 'auth':
+            auth_header = data.get('authorization', '')
+            token = auth_header.replace('Bearer ', '') if auth_header.startswith('Bearer ') else None
+
+            if await self._validate_token(token):
+                ws._authenticated = True
+                await ws.send_json({'type': 'auth_success', 'message': 'Authentication successful'})
+            else:
+                await ws.send_json({'type': 'auth_error', 'message': 'Invalid or expired token'})
+            return
+
+        if msg_type == 'stream_request':
+            prompt, ask_kwargs = self._extract_stream_params(data, 'type')
+
+            # Send acknowledgment
+            await ws.send_json({
+                'type': 'stream_start',
+                'prompt': prompt
+            })
+
+            try:
+                # Stream the LLM response
+                async for chunk in bot.ask_stream(prompt, **ask_kwargs):
+                    await ws.send_json({
+                        'type': 'content',
+                        'data': chunk
+                    })
+
+                # Send completion
+                await ws.send_json({
+                    'type': 'stream_complete'
+                })
+
+            except Exception as e:
+                await ws.send_json({
+                    'type': 'error',
+                    'message': str(e)
+                })
+
+        elif msg_type == 'ping':
+            await ws.send_json({'type': 'pong'})
+
+        else:
+            await ws.send_json({
+                'type': 'error',
+                'message': f'Unknown message type: {msg_type}'
+            })
+
+    async def broadcast(self, message: dict):
+        """Broadcast message to all connected clients"""
+        for ws in self.active_connections:
+            try:
+                await ws.send_json(message)
+            except Exception as e:
+                print(f"Error broadcasting to client: {e}")
+
+    def configure_routes(self, app: web.Application):
+        """Configure routes for streaming endpoints."""
+        # sse endpoint
+        exclude_list.append('/bots/*/stream/sse')
+        app.router.add_post('/bots/{bot_id}/stream/sse', self.stream_sse)
+        # ndjson endpoint
+        exclude_list.append('/bots/*/stream/ndjson')
+        app.router.add_post('/bots/{bot_id}/stream/ndjson', self.stream_ndjson)
+        # chunked endpoint
+        exclude_list.append('/bots/*/stream/chunked')
+        app.router.add_post('/bots/{bot_id}/stream/chunked', self.stream_chunked)
+        # websocket endpoint
+        exclude_list.append('/bots/*/stream/ws')
+        app.router.add_get('/bots/{bot_id}/stream/ws', self.stream_websocket)

parrot/interfaces/__init__.py

@@ -0,0 +1,6 @@
+from .database import DBInterface
+
+
+__all__ = (
+    'DBInterface',
+)

parrot/interfaces/aws.py

@@ -0,0 +1,143 @@
+"""
+AWS Interface for AI-Parrot
+Provides async context manager for AWS service clients using aioboto3
+"""
+from typing import Optional, Dict, Any, AsyncIterator
+from contextlib import asynccontextmanager
+import aioboto3
+from botocore.exceptions import ClientError, NoCredentialsError
+from ..conf import (
+    AWS_ACCESS_KEY,
+    AWS_SECRET_KEY,
+    AWS_REGION_NAME,
+    AWS_CREDENTIALS
+)
+
+
+class AWSInterface:
+    """
+    Base interface for AWS services using aioboto3.
+    
+    Provides async context manager for creating service clients.
+    Handles credential management and session lifecycle.
+    
+    Example:
+        >>> aws = AWSInterface(aws_id='default')
+        >>> async with aws.client('s3') as s3:
+        ...     response = await s3.list_buckets()
+    """
+    
+    def __init__(
+        self,
+        aws_id: str = 'default',
+        region_name: Optional[str] = None,
+        credentials: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ):
+        """
+        Initialize AWS interface.
+        
+        Args:
+            aws_id: Identifier for credentials in AWS_CREDENTIALS dict
+            region_name: AWS region (overrides credentials config)
+            credentials: Direct credential dict (overrides aws_id lookup)
+            **kwargs: Additional boto3 session parameters
+        """
+        # Get credentials from config or direct input
+        if credentials is None:
+            credentials = AWS_CREDENTIALS.get(aws_id, {})
+            if not credentials or credentials == 'default':
+                credentials = AWS_CREDENTIALS.get('default', {
+                    'aws_key': AWS_ACCESS_KEY,
+                    'aws_secret': AWS_SECRET_KEY,
+                    'region_name': AWS_REGION_NAME
+                })
+        
+        # Build AWS config
+        self.aws_config = {
+            'aws_access_key_id': credentials.get('aws_key') or credentials.get('aws_access_key_id'),
+            'aws_secret_access_key': credentials.get('aws_secret') or credentials.get('aws_secret_access_key'),
+            'region_name': region_name or credentials.get('region_name', AWS_REGION_NAME),
+        }
+        
+        # Add optional session token if present
+        if 'aws_session_token' in credentials:
+            self.aws_config['aws_session_token'] = credentials['aws_session_token']
+        
+        # Add any additional kwargs
+        self.aws_config.update(kwargs)
+        
+        # Remove None values
+        self.aws_config = {k: v for k, v in self.aws_config.items() if v is not None}
+        
+        # Create session
+        self.session = aioboto3.Session(**self.aws_config)
+        self._region = self.aws_config.get('region_name')
+    
+    @property
+    def region(self) -> str:
+        """Get configured AWS region"""
+        return self._region
+    
+    @asynccontextmanager
+    async def client(self, service_name: str, **kwargs) -> AsyncIterator[Any]:
+        """
+        Async context manager for AWS service client.
+        
+        Args:
+            service_name: AWS service name (e.g., 's3', 'cloudwatch', 'logs')
+            **kwargs: Additional client configuration
+            
+        Yields:
+            AWS service client
+            
+        Example:
+            >>> async with aws.client('cloudwatch') as cw:
+            ...     metrics = await cw.list_metrics()
+        """
+        async with self.session.client(service_name, **kwargs) as client:
+            yield client
+    
+    @asynccontextmanager
+    async def resource(self, service_name: str, **kwargs) -> AsyncIterator[Any]:
+        """
+        Async context manager for AWS service resource.
+        
+        Args:
+            service_name: AWS service name (e.g., 's3', 'dynamodb')
+            **kwargs: Additional resource configuration
+            
+        Yields:
+            AWS service resource
+        """
+        async with self.session.resource(service_name, **kwargs) as resource:
+            yield resource
+    
+    async def validate_credentials(self) -> bool:
+        """
+        Validate AWS credentials by making a simple API call.
+        
+        Returns:
+            True if credentials are valid, False otherwise
+        """
+        try:
+            async with self.client('sts') as sts:
+                await sts.get_caller_identity()
+            return True
+        except (ClientError, NoCredentialsError):
+            return False
+    
+    async def get_caller_identity(self) -> Dict[str, Any]:
+        """
+        Get AWS caller identity information.
+        
+        Returns:
+            Dict with UserId, Account, and Arn
+        """
+        async with self.client('sts') as sts:
+            response = await sts.get_caller_identity()
+            return {
+                'user_id': response.get('UserId'),
+                'account': response.get('Account'),
+                'arn': response.get('Arn')
+            }

parrot/interfaces/credentials.py

@@ -0,0 +1,113 @@
+import os
+from abc import ABC
+import contextlib
+from typing import TypeVar
+from typing_extensions import ParamSpec
+from navconfig import config
+from navconfig.logging import logging
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+valid_types = {
+    "<class 'str'>": str,
+    "<class 'int'>": int,
+    "<class 'float'>": float,
+    "<class 'list'>": list,
+    "<class 'tuple'>": tuple,
+    "<class 'dict'>": dict
+}
+
+
+class CredentialsInterface(ABC):
+    """
+    Abstract Base Class for handling credentials and environment variables.
+    This class provides methods to process and validate credentials, as well as
+    retrieve values from environment variables or configuration files.
+    """
+    _credentials: dict = {"username": str, "password": str}
+
+    def __init__(self, *args, **kwargs) -> None:
+        # if credentials:
+        self.credentials: dict = kwargs.pop('credentials', None)
+        self._no_warnings = kwargs.get("no_warnings", False)
+        if expected := kwargs.pop("expected_credentials", None):
+            self._credentials = expected
+        self._environment = config
+        try:
+            super().__init__(*args, **kwargs)
+        except TypeError:
+            super().__init__()
+        # Interface not started:
+        self._started: bool = False
+        self.logger = logging.getLogger(__name__)
+
+    def get_env_value(self, key, default: str = None, expected_type: object = None):
+        """
+        Retrieves a value from the environment variables or the configuration.
+
+        :param key: The key for the environment variable.
+        :param default: The default value to return if the key is not found.
+        :return: The value of the environment variable or the default value.
+        """
+        if key is None:
+            return default
+        if expected_type is not None:
+            with contextlib.suppress(TypeError):
+                if expected_type in (int, float):
+                    return val if (val := self._environment.getint(key)) else key
+                elif expected_type == bool:
+                    return val if (val := self._environment.getboolean(key)) else key
+                else:
+                    return val if (val := self._environment.get(key)) else key
+            return default
+        if val := os.getenv(str(key), default):
+            return val
+        return val if (val := self._environment.get(key, default)) else key
+
+    def processing_credentials(self):
+        if self.credentials:
+            for key, expected_type in self._credentials.items():
+                try:
+                    value = self.credentials.get(key, None)
+                    default = getattr(self, key, value)
+                    # print('KEY ', key, 'VAL ', value, 'DEF ', default)
+                    if type(value) == expected_type or isinstance(value, valid_types[str(expected_type)]):  # pylint: disable=E1136 # noqa
+                        # can process the credentials, extracted from environment or variables:
+                        val = self.get_env_value(
+                            value, default=default, expected_type=expected_type
+                        )
+                        # print('VAL > ', val, 'DEFAULT > ', default, expected_type)
+                        self.credentials[key] = val
+                        # print('KEY: ', key, self.credentials[key])
+                    elif isinstance(value, str):
+                        # Use os.getenv to get the value from environment variables
+                        env_value = self.get_env_value(
+                            value, default=default, expected_type=expected_type
+                        )
+                        self.credentials[key] = env_value
+                    else:
+                        self.credentials[key] = default
+                except KeyError as exc:
+                    print(f'Failed credential {key} with value {value}: {exc}')
+                    continue
+                except (TypeError, ValueError) as ex:
+                    self.logger.error(f"{__name__}: Wrong or missing Credentials")
+                    raise RuntimeError(
+                        f"{__name__}: Wrong or missing Credentials"
+                    ) from ex
+                except Exception as ex:
+                    self.logger.exception(
+                        f"Error Processing Credentials: {ex}"
+                    )
+                    raise RuntimeError(
+                        f"Error Processing Credentials: {ex}"
+                    ) from ex
+        if self.credentials is None:
+            if self._no_warnings is False:
+                self.logger.warning(
+                    "No credentials where Found."
+                )
+            self.credentials = {}

parrot/interfaces/database.py

@@ -0,0 +1,27 @@
+"""DB (asyncdb) Extension.
+DB connection for any Application.
+"""
+from asyncdb import AsyncDB
+
+
+class DBInterface:
+    """
+    Interface for using database connections in an Application using AsyncDB.
+    """
+
+    def get_database(
+        self,
+        driver: str,
+        dsn: str = None,
+        params: dict = None,
+        timeout: int = 60,
+        **kwargs
+    ) -> AsyncDB:
+        """Get the driver."""
+        return AsyncDB(
+            driver,
+            dsn=dsn,
+            params=params,
+            timeout=timeout,
+            **kwargs
+        )