rapida-python 0.1.0__py3-none-any.whl → 0.1.0a0__py3-none-any.whl

rapida/__init__.py

@@ -1,13 +1,38 @@
-"""
-author: prashant.srivastav
-"""
+#  Copyright (c) 2024. Rapida
+#
+#  Permission is hereby granted, free of charge, to any person obtaining a copy
+#  of this software and associated documentation files (the "Software"), to deal
+#  in the Software without restriction, including without limitation the rights
+#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#  copies of the Software, and to permit persons to whom the Software is
+#  furnished to do so, subject to the following conditions:
+#
+#  The above copyright notice and this permission notice shall be included in
+#  all copies or substantial portions of the Software.
+#
+#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+#  THE SOFTWARE.
+#
+#  Author: Prashant <prashant@rapida.ai>
 
-from rapida.rapida_client import RapidaClient
+
+from rapida.rapida_client import (
+    RapidaClient,
+    RapidaEndpointClient,
+    RapidaAssistantClient,
+    RapidaGatewayClient,
+)
 from rapida.rapida_client_options import (
     RapidaClientOptions,
-    RapidaEnvironment,
-    RapidaRegion,
 )
+from rapida.rapida_environment import RapidaEnvironment
+from rapida.rapida_region import RapidaRegion
+from rapida.rapida_source import RapidaSource
 from rapida.exceptions import (
     RapidaException,
     RapidaInternalServerException,
@@ -15,9 +40,14 @@ from rapida.exceptions import (
     RapidaConfigurationException,
 )
 from rapida.version import VERSION
+from rapida.client.response_wrapper import Message, Content, ToolDefinition
 
 __all__ = [
     "RapidaClient",
+    "ToolDefinition",
+    "RapidaEndpointClient",
+    "RapidaAssistantClient",
+    "RapidaGatewayClient",
     "RapidaClientOptions",
     "RapidaException",
     "RapidaInternalServerException",
@@ -26,4 +56,7 @@ __all__ = [
     "VERSION",
     "RapidaEnvironment",
     "RapidaRegion",
+    "RapidaSource",
+    "Message",
+    "Content",
 ]

rapida/client/grpc_bridge.py

@@ -1,33 +1,56 @@
-"""
-author: prashant.srivastav
+#  Copyright (c) 2024. Rapida
+#
+#  Permission is hereby granted, free of charge, to any person obtaining a copy
+#  of this software and associated documentation files (the "Software"), to deal
+#  in the Software without restriction, including without limitation the rights
+#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#  copies of the Software, and to permit persons to whom the Software is
+#  furnished to do so, subject to the following conditions:
+#
+#  The above copyright notice and this permission notice shall be included in
+#  all copies or substantial portions of the Software.
+#
+#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+#  THE SOFTWARE.
+#
+# Base Bridge class for all RPC and gRPC clients.
+# - Changing this class will impact all RPC and gRPC clients and consumer service which implement and inherit
+#     this class.
+# - Considered as core computability layer for all RPC and gRPC clients the CHANGE REQUEST is not allowed.
+# - Mostly response being handled with JSON format.
+#
+#
+#  Author: Prashant <prashant@rapida.ai>
 
-Base Bridge class for all RPC and gRPC clients.
-- Changing this class will impact all RPC and gRPC clients and consumer service which implement and inherit
-    this class.
-- Considered as core computability layer for all RPC and gRPC clients the CHANGE REQUEST is not allowed.
-- Mostly response being handled with JSON format.
 
-"""
 import logging
 import time
 from abc import ABC
-from typing import Any, Dict
+from typing import Any, AsyncIterator, Dict
 
 import grpc
 from google.protobuf.json_format import MessageToDict
-from grpc import aio as grpc_aio
+from grpc import UnaryStreamMultiCallable, aio as grpc_aio
 from google.protobuf.message import Message
 from grpc.aio import Metadata
 from grpc.aio._channel import UnaryUnaryMultiCallable
 
+from rapida.constants import (
+    HEADER_API_KEY,
+    HEADER_ENVIRONMENT_KEY,
+    HEADER_REGION_KEY,
+    HEADER_SOURCE_KEY,
+)
 from rapida.exceptions.exceptions import RapidaException
 
 _log = logging.getLogger("rapida.client.grpc_bridge")
 
 
-# _log = logging.getLogger("bridges.grpc_bridge")
-
-
 class GRPCBridge(ABC):
     """
     Base class for all GRPC bridges.
@@ -38,6 +61,7 @@ class GRPCBridge(ABC):
     rapida_api_key: str
     rapida_region: str
     rapida_environment: str
+    rapida_source: str
     rapida_is_secure: bool
 
     @classmethod
@@ -47,6 +71,7 @@ class GRPCBridge(ABC):
         rapida_api_key: str,
         rapida_region: str,
         rapida_environment: str,
+        rapida_source: str,
         rapida_is_secure: bool,
     ):
         """
@@ -54,11 +79,13 @@ class GRPCBridge(ABC):
             service_url: a url where the endpoint service is deployed
             rapida_api_key: a api keys to authenticate and authorize the request made to rapida servers (Please find the docs <a>docs.rapida.ai</a>
             rapida_region: specify the region to accelerate the request if it's specific as per client need
+            rapida_source: source of the request for rapida request later it will become important for routing and parsing the information
             rapida_environment: a tag to represent the env currently (production and development is supported)
         """
         cls.service_url = service_url
         cls.rapida_api_key = rapida_api_key
         cls.rapida_region = rapida_region
+        cls.rapida_source = rapida_source
         cls.rapida_environment = rapida_environment
         cls.rapida_is_secure = rapida_is_secure
 
@@ -108,9 +135,10 @@ class GRPCBridge(ABC):
                 # metadata for request
                 _metadata: Metadata = Metadata()
 
-                _metadata.add("x-api-key", cls.rapida_api_key)
-                _metadata.add("x-rapida-environment", cls.rapida_environment)
-                _metadata.add("x-rapida-region", cls.rapida_region)
+                _metadata.add(HEADER_API_KEY, cls.rapida_api_key)
+                _metadata.add(HEADER_ENVIRONMENT_KEY, cls.rapida_environment)
+                _metadata.add(HEADER_REGION_KEY, cls.rapida_region)
+                _metadata.add(HEADER_SOURCE_KEY, cls.rapida_source)
 
                 # request executed with asynchronous invocation of a unary-call RPC.
                 results = await unary_unary(request=message_type, metadata=_metadata)
@@ -126,3 +154,52 @@ class GRPCBridge(ABC):
                 f"{cls.__qualname__} {attr} [status:NOT_OK request:{time.time() - started_request}s]"
             )
             raise RapidaException(message=str(ex), code=500, source=__name__)
+
+    @classmethod
+    async def fetch_stream(
+        cls,
+        stub: Any,
+        attr: str,
+        message_type: Message,  # Unary request
+        **unmarshal_options,
+    ) -> AsyncIterator[Dict[str, Any]]:
+        """
+        Generic requestor for unary request with streaming response
+        :param stub:
+        :param attr:
+        :param message_type: Unary request message
+        :param unmarshal_options: Options to control unmarshalling of the protobuf result.
+        """
+        started_request = time.time()
+        try:
+            _log.debug(f"grpc_bridge: in span of {attr}")
+            async with cls.channel() as channel:
+                _log.debug(f"grpc_bridge: created secure channel for {cls.service_url}")
+
+                # Get unary-stream method from the stub (request once, receive a stream)
+                unary_stream: UnaryStreamMultiCallable = getattr(stub(channel), attr)
+
+                # Metadata for the request
+                _metadata: Metadata = Metadata()
+                _metadata.add(HEADER_API_KEY, cls.rapida_api_key)
+                _metadata.add(HEADER_ENVIRONMENT_KEY, cls.rapida_environment)
+                _metadata.add(HEADER_REGION_KEY, cls.rapida_region)
+                _metadata.add(HEADER_SOURCE_KEY, cls.rapida_source)
+
+                # Send the unary request and receive a stream of responses
+                async for response in unary_stream(
+                    request=message_type, metadata=_metadata
+                ):
+                    # Unmarshal each result with the given options
+                    json_result = MessageToDict(response, **unmarshal_options)
+                    _log.info(
+                        f"{cls.__qualname__} {attr} [status:OK request:{time.time() - started_request}s]"
+                    )
+                    # Yield each response as part of the stream
+                    yield json_result
+        except grpc_aio.AioRpcError as ex:
+            # Handle errors from inter-service client communication
+            _log.error(
+                f"{cls.__qualname__} {attr} [status:NOT_OK request:{time.time() - started_request}s]"
+            )
+            raise RapidaException(message=str(ex), code=500, source=__name__)

rapida/client/rapida_bridge.py

@@ -1,14 +1,36 @@
-"""
-author: prashant.srivastav
-"""
+#  Copyright (c) 2024. Rapida
+#
+#  Permission is hereby granted, free of charge, to any person obtaining a copy
+#  of this software and associated documentation files (the "Software"), to deal
+#  in the Software without restriction, including without limitation the rights
+#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#  copies of the Software, and to permit persons to whom the Software is
+#  furnished to do so, subject to the following conditions:
+#
+#  The above copyright notice and this permission notice shall be included in
+#  all copies or substantial portions of the Software.
+#
+#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+#  THE SOFTWARE.
+#
+#  Author: Prashant <prashant@rapida.ai>
+
+
 import logging
-from typing import Dict, Optional, Mapping
+from typing import AsyncIterator, Dict, List, Optional, Mapping
 from google.protobuf.json_format import ParseDict
 from google.protobuf.any_pb2 import Any
-
-
+from google.protobuf.struct_pb2 import Struct
 from rapida.client.response_wrapper import InvokeResponseWrapper
 from rapida.artifacts.protos import (
+    common_pb2,
+    integration_api_pb2,
+    integration_api_pb2_grpc,
     invoker_api_pb2,
     invoker_api_pb2_grpc,
 )
@@ -24,6 +46,7 @@ class RapidaBridge(GRPCBridge):
         rapida_api_key: str,
         rapida_region: str,
         rapida_environment: str,
+        rapida_source: str,
         rapida_is_secure: bool,
     ):
         """
@@ -39,6 +62,7 @@ class RapidaBridge(GRPCBridge):
             rapida_api_key,
             rapida_region,
             rapida_environment,
+            rapida_source,
             rapida_is_secure,
         )
 
@@ -96,7 +120,6 @@ class RapidaBridge(GRPCBridge):
             attr="Probe",
             message_type=invoker_api_pb2.ProbeRequest(requestId=rapida_audit_id),
             preserving_proto_field_name=True,
-            # including_default_value_fields=True,
         )
         return ParseDict(
             response, invoker_api_pb2.ProbeResponse(), ignore_unknown_fields=True
@@ -116,7 +139,7 @@ class RapidaBridge(GRPCBridge):
         """
         response = await self.fetch(
             stub=invoker_api_pb2_grpc.DeploymentStub,
-            attr="UpdateMetadata",
+            attr="Update",
             message_type=invoker_api_pb2.UpdateRequest(
                 requestId=rapida_audit_id,
                 metadata=rapida_metadata,
@@ -128,3 +151,239 @@ class RapidaBridge(GRPCBridge):
         return ParseDict(
             response, invoker_api_pb2.UpdateResponse(), ignore_unknown_fields=True
         )
+
+    async def chat(
+        self,
+        cred: Dict,
+        provider: str,
+        model: str,
+        conversations: List[common_pb2.Message],
+        tool_definitions: List[integration_api_pb2.ToolDefinition],
+        model_parameters: List[integration_api_pb2.ModelParameter],
+        meta: Mapping[str, str] = None,
+    ) -> integration_api_pb2.ChatResponse:
+        """
+        Asynchronously initiates a chat interaction with the Rapida service.
+
+        This method sends a chat request to the Rapida service and processes the response.
+        It utilizes gRPC to communicate with the service, specifically calling the
+        'UpdateMetadata' method of the RapidaServiceStub.
+
+        The function performs the following steps:
+        1. Sends a chat request to the Rapida service using the 'fetch' method.
+        2. Waits for the asynchronous response from the service.
+        3. Parses the received response into a ChatResponse object.
+
+        Returns:
+            integration_api_pb2.ChatResponse: The processed response from the Rapida service,
+            containing the chat interaction results.
+
+        Note:
+            - This method uses the 'fetch' utility to handle the gRPC communication.
+            - The 'preserving_proto_field_name' parameter is set to True to maintain
+            the original field names in the protocol buffer message.
+            - The response is parsed using the 'ParseDict' function, which converts
+            the dictionary response into a ChatResponse protocol buffer object.
+            - Unknown fields in the response are ignored during parsing.
+
+        Raises:
+            Any exceptions that might occur during the gRPC communication or response parsing.
+
+        Example usage:
+            chat_response = await rapida_bridge_instance.chat()
+            # Process the chat_response as needed
+        """
+
+        credentials = Struct()
+        for k, v in cred.items():
+            credentials.update({k: v})
+
+        response = await self.fetch(
+            stub=integration_api_pb2_grpc.RapidaServiceStub,
+            attr="Chat",
+            message_type=integration_api_pb2.ChatRequest(
+                credential=integration_api_pb2.Credential(value=credentials),
+                model=f"{provider}/{model}",
+                conversations=conversations,
+                toolDefinitions=tool_definitions,
+                modelParameters=model_parameters,
+                additionalData=meta,
+            ),
+            preserving_proto_field_name=True,
+        )
+
+        return ParseDict(
+            response, integration_api_pb2.ChatResponse(), ignore_unknown_fields=True
+        )
+
+    async def generate(
+        self,
+        cred: Dict,
+        provider: str,
+        model: str,
+        system_prompt: str,
+        prompt: str,
+        model_parameters: Mapping[str, str] = None,
+        meta: Mapping[str, str] = None,
+    ) -> integration_api_pb2.GenerateResponse:
+        """
+        Prepares and sends a generation request to the Rapida service.
+
+        This function:
+        1. Converts credentials and model parameters to the required format
+        2. Constructs a GenerateRequest message with the provided inputs
+        3. Sends the request to the Rapida service using gRPC
+        4. Parses and returns the response as a GenerateResponse object
+
+        Args:
+            cred (dict): Credentials for authentication
+            system_prompt (str): System prompt for the generation
+            prompt (str): User prompt for the generation
+            provider (str): The AI provider to use
+            model (str): The specific model to use
+            model_parameters (dict): Additional parameters for the model
+            meta (dict): Any additional metadata
+
+        Returns:
+            integration_api_pb2.GenerateResponse: The parsed response from the service
+        """
+        credentials = Struct()
+        for k, v in cred.items():
+            credentials.update({k: v})
+
+        args: List[integration_api_pb2.ModelParameter] = []
+        for k, v in model_parameters.items():
+            args.append(
+                integration_api_pb2.ModelParameter(
+                    key=k,
+                    value=v,
+                )
+            )
+
+        response = await self.fetch(
+            stub=integration_api_pb2_grpc.RapidaServiceStub,
+            attr="Generate",
+            message_type=integration_api_pb2.GenerateRequest(
+                credential=credentials,
+                systemPrompt=system_prompt,
+                prompt=prompt,
+                model=f"{provider}/{model}",
+                modelParameters=args,
+                additionalData=meta,
+            ),
+            preserving_proto_field_name=True,
+        )
+
+        return ParseDict(
+            response, integration_api_pb2.GenerateResponse(), ignore_unknown_fields=True
+        )
+
+    async def embedding(
+        self,
+        cred: Dict,
+        provider: str,
+        model: str,
+        contents: List[str],
+        model_parameters: Mapping[str, str] = None,
+        meta: Mapping[str, str] = None,
+    ) -> integration_api_pb2.EmbeddingResponse:
+
+        credentials = Struct()
+        for k, v in cred.items():
+            credentials.update({k: v})
+
+        args: List[integration_api_pb2.ModelParameter] = []
+        for k, v in model_parameters.items():
+            args.append(
+                integration_api_pb2.ModelParameter(
+                    key=k,
+                    value=v,
+                )
+            )
+
+        _content: Mapping[int, str] = {}
+        for i, val in enumerate(contents):
+            _content[i] = val
+
+        response = await self.fetch(
+            stub=integration_api_pb2_grpc.RapidaServiceStub,
+            attr="Embedding",
+            message_type=integration_api_pb2.EmbeddingRequest(
+                credential=credentials,
+                model=f"{provider}/{model}",
+                modelParameters=args,
+                additionalData=meta,
+                content=_content,
+            ),
+            preserving_proto_field_name=True,
+        )
+
+        return ParseDict(
+            response,
+            integration_api_pb2.EmbeddingResponse(),
+            ignore_unknown_fields=True,
+        )
+
+    async def chat_stream(
+        self,
+        cred: Dict,
+        provider: str,
+        model: str,
+        conversations: List[common_pb2.Message],
+        tool_definitions: List[integration_api_pb2.ToolDefinition],
+        model_parameters: List[integration_api_pb2.ModelParameter],
+        meta: Mapping[str, str] = None,
+    ) -> AsyncIterator[integration_api_pb2.ChatResponse]:
+        """
+        Asynchronously initiates a chat interaction with the Rapida service and streams the responses.
+
+        This method sends a chat request to the Rapida service and handles streaming responses.
+        It utilizes gRPC to communicate with the service, specifically calling the
+        'StreamChat' method of the RapidaServiceStub.
+
+        The function performs the following steps:
+        1. Sends a chat request to the Rapida service using the 'fetch_stream_response' method.
+        2. Asynchronously iterates over the streaming responses from the service.
+        3. Parses each received response into a ChatResponse object.
+
+        Returns:
+            AsyncIterator[integration_api_pb2.ChatResponse]: The streamed responses from the Rapida service,
+            containing the chat interaction results.
+
+        Note:
+            - This method uses the 'fetch_stream' utility to handle the gRPC communication.
+            - The 'preserving_proto_field_name' parameter is set to True to maintain
+            the original field names in the protocol buffer message.
+            - Each response is parsed using the 'ParseDict' function, which converts
+            the dictionary response into a ChatResponse protocol buffer object.
+            - Unknown fields in the response are ignored during parsing.
+
+        Raises:
+            Any exceptions that might occur during the gRPC communication or response parsing.
+
+        Example usage:
+            async for chat_response in rapida_bridge_instance.chat_stream():
+                # Process each chat_response as needed
+        """
+
+        credentials = Struct()
+        for k, v in cred.items():
+            credentials.update({k: v})
+
+        # Stream the chat response from the gRPC server
+        async for response in self.fetch_stream(
+            stub=integration_api_pb2_grpc.RapidaServiceStub,
+            attr="StreamChat",
+            message_type=integration_api_pb2.ChatRequest(
+                credential=integration_api_pb2.Credential(value=credentials),
+                model=f"{provider}/{model}",
+                conversations=conversations,
+                toolDefinitions=tool_definitions,
+                modelParameters=model_parameters,
+                additionalData=meta,
+            ),
+            preserving_proto_field_name=True,
+        ):
+            yield ParseDict(
+                response, integration_api_pb2.ChatResponse(), ignore_unknown_fields=True
+            )