├── layers └── openai │ └── requirements.txt ├── images └── openai.png ├── lambdas ├── connect_lambda │ └── lambda_function.py ├── disconnect_lambda │ └── lambda_function.py └── streaming_lambda │ ├── lambda_function.py │ └── invoke_openai.py ├── client └── client.py ├── README.md └── template.yaml /layers/openai/requirements.txt: -------------------------------------------------------------------------------- 1 | openai==0.27.8 2 | urllib3==1.26.6 3 | -------------------------------------------------------------------------------- /images/openai.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rp-86/streaming_openai_aws/HEAD/images/openai.png -------------------------------------------------------------------------------- /lambdas/connect_lambda/lambda_function.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | def lambda_handler(event, context): 4 | return { 5 | "statusCode": 200, 6 | "body": "Success" 7 | } -------------------------------------------------------------------------------- /lambdas/disconnect_lambda/lambda_function.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | def lambda_handler(event, context): 4 | return { 5 | "statusCode": 200, 6 | "body": "Success" 7 | } -------------------------------------------------------------------------------- /lambdas/streaming_lambda/lambda_function.py: -------------------------------------------------------------------------------- 1 | import json 2 | from invoke_openai import InvokeOpenai 3 | 4 | def lambda_handler(event, context): 5 | body = json.loads(event["body"]) 6 | request = body["query"] 7 | connectionId = event["requestContext"]["connectionId"] 8 | ivk = InvokeOpenai(connectionId) 9 | ivk.call_openai(request) 10 | 11 | return { 12 | "statusCode": 200, 13 | "body": "Success" 14 | } -------------------------------------------------------------------------------- /client/client.py: -------------------------------------------------------------------------------- 1 | import websocket 2 | import ssl 3 | import json 4 | 5 | def on_message(ws, message): 6 | print("Received message:", message) 7 | 8 | def on_error(ws, error): 9 | print("WebSocket error:", error) 10 | 11 | def on_close(ws, a, b): 12 | print("WebSocket connection closed") 13 | 14 | def on_open(ws): 15 | print("WebSocket connection opened") 16 | # Send a request after the connection is opened 17 | request = { 18 | "action": "openai", 19 | "query":"What is AI?" 20 | } 21 | ws.send(json.dumps(request)) 22 | 23 | if __name__ == "__main__": 24 | # Specify your WebSocket API Gateway endpoint URL 25 | websocket_url = "wss://er7uagkgg4.execute-api.ap-south-1.amazonaws.com/dev" 26 | 27 | # Create a WebSocket connection 28 | websocket.enableTrace(True) 29 | ws = websocket.WebSocketApp(websocket_url, 30 | on_message=on_message, 31 | on_error=on_error, 32 | on_close=on_close) 33 | 34 | # Use SSL/TLS for secure WebSocket connection 35 | ws.on_open = on_open 36 | ws.run_forever() 37 | -------------------------------------------------------------------------------- /lambdas/streaming_lambda/invoke_openai.py: -------------------------------------------------------------------------------- 1 | import boto3 2 | import os 3 | import openai 4 | 5 | class InvokeOpenai: 6 | def __init__(self, connectionId): 7 | self.ssm = boto3.client("ssm") 8 | self.conn = boto3.client("apigatewaymanagementapi", endpoint_url=os.environ["api_endpoint"], region_name=os.environ["region"]) 9 | self.params = { 10 | "Data":"", 11 | "ConnectionId": connectionId 12 | } 13 | 14 | def read_ssm_parameter(self): 15 | openai_key = self.ssm.get_parameter(Name=os.environ["openai_key"],WithDecryption=True)["Parameter"]["Value"] 16 | return openai_key 17 | 18 | def call_openai(self, request, model="gpt-3.5-turbo"): 19 | openai.api_key = self.read_ssm_parameter() 20 | response = "" 21 | 22 | for resp in openai.ChatCompletion.create( 23 | model=model, 24 | messages=[ 25 | {"role": "system", "content": "You are a helpful assistant."}, 26 | {"role": "user", "content": f"{request}"} 27 | ], 28 | stream=True, 29 | stop=None 30 | ): 31 | if "content" in resp.choices[0]["delta"]: 32 | res = resp.choices[0]["delta"]["content"] 33 | response += res 34 | if res != '': 35 | self.params["Data"] = res 36 | self.conn.post_to_connection(**self.params) 37 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Real-Time AI-Powered WebSocket APIs 2 | This repository provides a demonstration of building real-time AI-powered WebSocket APIs using AWS API Gateway, Lambda, and OpenAI. The setup enables bidirectional communication between clients and the backend, with AI processing performed by OpenAI. 3 | 4 | ## Architecture Overview 5 | 6 | ![Architecture](images/openai.png) 7 | 8 | The architecture consists of the following components: 9 | 10 | - AWS WebSocket API Gateway: Serves as the entry point for WebSocket communication, facilitating real-time, bidirectional communication between clients and the backend. 11 | - AWS Lambda Functions: 12 | - Connect Lambda: Handles the initial connection request and establishes the WebSocket connection. 13 | - Streaming OpenAI Lambda: Processes client requests, calls the OpenAI API, and sends AI-powered responses back to clients. 14 | - OpenAI API: Provides AI capabilities and generates responses based on client queries. 15 | 16 | ## Getting Started 17 | To get started with the project, follow the steps outlined below: 18 | 19 | ### Prerequisites 20 | An AWS account 21 | Basic knowledge of AWS services, including API Gateway, Lambda, and CloudFormation 22 | An OpenAI API key (sign up on the OpenAI website if you don't have one) 23 | ### Installation and Setup 24 | 1. Clone this repository to your local machine. 25 | 2. Configure the AWS CLI with your AWS account credentials. 26 | 3. Modify the SAM template (template.yaml) with your desired configurations, such as function names, IAM roles, and OpenAI API key storage in AWS SSM Parameter Store. 27 | 4. Deploy the SAM template using AWS CloudFormation. This will provision the necessary AWS resources, including the API Gateway and Lambda functions. 28 | 5. Configure the client-side script (client.py) by providing the WebSocket API Gateway URL. 29 | 6. Run the client-side script locally to test the WebSocket connection and observe the AI-powered responses. 30 | ### Additional Customizations 31 | Feel free to explore and customize the project further based on your requirements: 32 | 33 | - Modify the Lambda function code to enhance or modify the AI processing logic. 34 | - Implement additional functionalities, such as error handling or security measures. 35 | - Integrate with other AI services or APIs to extend the capabilities of the WebSocket API. 36 | ## Contributing 37 | Contributions to this project are welcome! If you have any ideas, improvements, or bug fixes, please open an issue or submit a pull request. 38 | 39 | ## Resources 40 | - AWS Serverless Application Model (SAM) Documentation 41 | - AWS CloudFormation Documentation 42 | - AWS Lambda Documentation 43 | - OpenAI API Documentation 44 | - WebSocket API Gateway Documentation 45 | 46 | Please refer to the above resources for detailed documentation on AWS services and the OpenAI API. 47 | 48 | ## Acknowledgments 49 | This project was inspired by the need for real-time AI-powered communication and the capabilities offered by AWS services and OpenAI. 50 | We would like to express our gratitude to the developers and contributors of the AWS SDKs, OpenAI Python library, and related open-source projects that facilitated the implementation of this project. 51 | 52 | 53 | 54 | -------------------------------------------------------------------------------- /template.yaml: -------------------------------------------------------------------------------- 1 | AWSTemplateFormatVersion: '2010-09-09' 2 | Transform: 'AWS::Serverless-2016-10-31' 3 | 4 | 5 | 6 | Resources: 7 | MyIAMRole: 8 | Type: AWS::IAM::Role 9 | Properties: 10 | RoleName: service-role 11 | AssumeRolePolicyDocument: 12 | Version: '2012-10-17' 13 | Statement: 14 | - Effect: Allow 15 | Principal: 16 | Service: lambda.amazonaws.com 17 | Action: sts:AssumeRole 18 | Policies: 19 | - PolicyName: CloudWatchLogsFullAccess 20 | PolicyDocument: 21 | Version: '2012-10-17' 22 | Statement: 23 | - Effect: Allow 24 | Action: logs:* 25 | Resource: '*' 26 | - PolicyName: AmazonDynamoDBFullAccess 27 | PolicyDocument: 28 | Version: '2012-10-17' 29 | Statement: 30 | - Effect: Allow 31 | Action: dynamodb:* 32 | Resource: '*' 33 | - PolicyName: AmazonS3FullAccess 34 | PolicyDocument: 35 | Version: '2012-10-17' 36 | Statement: 37 | - Effect: Allow 38 | Action: s3:* 39 | Resource: '*' 40 | - PolicyName: AmazonSSMFullAccess 41 | PolicyDocument: 42 | Version: '2012-10-17' 43 | Statement: 44 | - Effect: Allow 45 | Action: ssm:* 46 | Resource: '*' 47 | - PolicyName: AmazonAPIGatewayInvokeFullAccess 48 | PolicyDocument: 49 | Version: '2012-10-17' 50 | Statement: 51 | - Effect: Allow 52 | Action: apigateway:* 53 | Resource: '*' 54 | - PolicyName: AWSLambda_FullAccess 55 | PolicyDocument: 56 | Version: '2012-10-17' 57 | Statement: 58 | - Effect: Allow 59 | Action: lambda:* 60 | Resource: '*' 61 | 62 | socketAPI: 63 | Type: AWS::ApiGatewayV2::Api 64 | Properties: 65 | Name: "socketAPI" 66 | ProtocolType: WEBSOCKET 67 | RouteSelectionExpression: "$request.body.action" 68 | 69 | ConnectInteg: 70 | Type: AWS::ApiGatewayV2::Integration 71 | DependsOn: MyIAMRole 72 | Properties: 73 | ApiId: !Ref socketAPI 74 | IntegrationType: AWS_PROXY 75 | CredentialsArn: !Sub 76 | - "{role}" 77 | - role: !GetAtt MyIAMRole.Arn 78 | IntegrationUri: !Sub 79 | - "arn:${AWS::Partition}:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${ConnectLambdaARN}/invocations" 80 | - ConnectLambdaARN: !GetAtt ConnectLambda.Arn 81 | 82 | ConnectRoute: 83 | Type: AWS::ApiGatewayV2::Route 84 | Properties: 85 | RouteKey: $connect 86 | ApiId: !Ref socketAPI 87 | OperationName: ConnectRoute 88 | Target: !Sub "integrations/${ConnectInteg}" 89 | 90 | ConnectLambda: 91 | Type: 'AWS::Serverless::Function' 92 | DependsOn: MyIAMRole 93 | Properties: 94 | FunctionName: 'connect-lambda' 95 | CodeUri: lambdas/connect_lambda/ 96 | Role: !Sub 97 | - "{role}" 98 | - role: !GetAtt MyIAMRole.Arn 99 | Handler: lambda_function.lambda_handler 100 | Runtime: python3.9 101 | Timeout: 600 102 | 103 | ConnectLambdaPermission: 104 | Type: AWS::Lambda::Permission 105 | DependsOn: 106 | - socketAPI 107 | Properties: 108 | FunctionName: !Ref ConnectLambda 109 | Action: 'lambda:InvokeFunction' 110 | Principal: apigateway.amazonaws.com 111 | SourceAccount: !Sub '${AWS::AccountId}' 112 | SourceArn: !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${socketAPI}/*/$connect' 113 | 114 | StreamingOpenai: 115 | Type: 'AWS::Serverless::Function' 116 | DependsOn: MyIAMRole 117 | Properties: 118 | FunctionName: 'streaming-openai' 119 | CodeUri: lambdas/streaming_lambda/ 120 | Role: !Sub 121 | - "{role}" 122 | - role: !GetAtt MyIAMRole.Arn 123 | Handler: lambda_function.lambda_handler 124 | Runtime: python3.9 125 | Timeout: 600 126 | Layers: 127 | - !Ref OpenaiLayer 128 | Environment: 129 | Variables: 130 | openai_key: "openai_key" 131 | region: !Sub "${AWS::Region}" 132 | api_endpoint: !Sub "https://${socketAPI}.execute-api.${AWS::Region}.amazonaws.com/dev" 133 | 134 | OpenaiInteg: 135 | Type: AWS::ApiGatewayV2::Integration 136 | DependsOn: MyIAMRole 137 | Properties: 138 | ApiId: !Ref socketAPI 139 | IntegrationType: AWS_PROXY 140 | CredentialsArn: !Sub 141 | - "{role}" 142 | - role: !GetAtt MyIAMRole.Arn 143 | IntegrationUri: !Sub 144 | - "arn:${AWS::Partition}:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${ConnectLambdaARN}/invocations" 145 | - ConnectLambdaARN: !GetAtt StreamingOpenai.Arn 146 | 147 | OpenaiRoute: 148 | Type: AWS::ApiGatewayV2::Route 149 | Properties: 150 | RouteKey: openai 151 | ApiId: !Ref socketAPI 152 | OperationName: OpenaiRoute 153 | Target: !Sub "integrations/${OpenaiInteg}" 154 | 155 | StreamingLambdaPermission: 156 | Type: AWS::Lambda::Permission 157 | DependsOn: 158 | - socketAPI 159 | Properties: 160 | FunctionName: !Ref StreamingOpenai 161 | Action: 'lambda:InvokeFunction' 162 | Principal: apigateway.amazonaws.com 163 | SourceAccount: !Sub '${AWS::AccountId}' 164 | SourceArn: !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${socketAPI}/*/openai' 165 | 166 | DisconnectInteg: 167 | Type: AWS::ApiGatewayV2::Integration 168 | DependsOn: MyIAMRole 169 | Properties: 170 | ApiId: !Ref socketAPI 171 | IntegrationType: AWS_PROXY 172 | CredentialsArn: !Sub 173 | - "{role}" 174 | - role: !GetAtt MyIAMRole.Arn 175 | IntegrationUri: !Sub 176 | - "arn:${AWS::Partition}:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${ConnectLambdaARN}/invocations" 177 | - ConnectLambdaARN: !GetAtt DiconnectLambda.Arn 178 | 179 | DisconnectRoute: 180 | Type: AWS::ApiGatewayV2::Route 181 | Properties: 182 | RouteKey: $disconnect 183 | ApiId: !Ref socketAPI 184 | OperationName: DisconnectRoute 185 | Target: !Sub "integrations/${DisconnectInteg}" 186 | 187 | DiconnectLambda: 188 | Type: 'AWS::Serverless::Function' 189 | DependsOn: MyIAMRole 190 | Properties: 191 | FunctionName: 'disconnect-lambda' 192 | CodeUri: lambdas/disconnect_lambda/ 193 | Role: !Sub 194 | - "{role}" 195 | - role: !GetAtt MyIAMRole.Arn 196 | Handler: lambda_function.lambda_handler 197 | Runtime: python3.9 198 | Timeout: 600 199 | 200 | DisconnectLambdaPermission: 201 | Type: AWS::Lambda::Permission 202 | DependsOn: 203 | - socketAPI 204 | Properties: 205 | FunctionName: !Ref DiconnectLambda 206 | Action: 'lambda:InvokeFunction' 207 | Principal: apigateway.amazonaws.com 208 | SourceAccount: !Sub '${AWS::AccountId}' 209 | SourceArn: !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${socketAPI}/*/$connect' 210 | 211 | 212 | OpenaiLayer: 213 | Type: AWS::Serverless::LayerVersion 214 | Properties: 215 | LayerName: 'openai-layer' 216 | Description: Layer for openai 217 | ContentUri: layers/openai/ 218 | CompatibleRuntimes: 219 | - python3.9 220 | Metadata: 221 | BuildMethod: python3.9 222 | 223 | Stage: 224 | Type: AWS::ApiGatewayV2::Stage 225 | Properties: 226 | StageName: dev 227 | Description: dev Stage 228 | ApiId: !Ref socketAPI 229 | 230 | Deployment: 231 | Type: AWS::ApiGatewayV2::Deployment 232 | DependsOn: 233 | - ConnectRoute 234 | - OpenaiRoute 235 | Properties: 236 | ApiId: !Ref socketAPI 237 | 238 | 239 | Outputs: 240 | socketAPI: 241 | Description: URL for the API Gateway 242 | Value: !Ref socketAPI --------------------------------------------------------------------------------