Skip to content
@aws/nx-plugin
Blog

AI Dungeon Game

Module 3: Story API implementation

The StoryApi comprises of a single API generate_story which given a Game and a list of Action’s for context, will progress a story. This API will be implemented as a streaming API in Python/FastAPI and will additionally demonstrate how changes can be made to the generated code to be fit for purpose.

API implementation

To create our API, we first need to install a couple of additional dependencies.

  • boto3 will be used to call Amazon Bedrock;
  • uvicorn will be used to start our API when used in conjunction with the Lambda Web Adapter (LWA).
  • copyfiles is an npm dependency that we will need to support cross-platform copying of files when updating our bundle task.

To install these dependencies, run the following commands:

Terminal window
pnpm nx run dungeon_adventure.story_api:add --args boto3 uvicorn
Terminal window
pnpm add -Dw copyfiles

Now let’s replace the contents of packages/story_api/story_api/main.py as follows:

packages/story_api/story_api/main.py
import json


from boto3 import client
from fastapi.responses import PlainTextResponse, StreamingResponse
from pydantic import BaseModel


from .init import app, lambda_handler


handler = lambda_handler


bedrock = client('bedrock-runtime')


class Action(BaseModel):
    role: str
    content: str


class StoryRequest(BaseModel):
    genre: str
    playerName: str
    actions: list[Action]


async def bedrock_stream(request: StoryRequest):
    messages = [
        {"role": "user", "content": "Continue or create a new story..."}
    ]


    for action in request.actions:
        messages.append({"role": action.role, "content": action.content})


    response = bedrock.invoke_model_with_response_stream(
        modelId='anthropic.claude-3-sonnet-20240229-v1:0',
        body=json.dumps({
            "system":f"""
            You are running an AI text adventure game in the {request.genre} genre.
            Player: {request.playerName}. Return less than 200 characters of text.
            """,
            "messages": messages,
            "max_tokens": 1000,
            "temperature": 0.7,
            "anthropic_version": "bedrock-2023-05-31"
        })
    )


    stream = response.get('body')
    if stream:
        for event in stream:
            chunk = event.get('chunk')
            if chunk:
                message = json.loads(chunk.get("bytes").decode())
                if message['type'] == "content_block_delta":
                    yield message['delta']['text'] or ""
                elif message['type'] == "message_stop":
                    yield "\n"


@app.post("/story/generate",
          openapi_extra={'x-streaming': True, 'x-query': True},
          response_class=PlainTextResponse)
def generate_story(request: StoryRequest) -> str:
    return StreamingResponse(bedrock_stream(request), media_type="text/plain")

Analyzing the code above:

  • We use the x-streaming setting to indicate that this is a streaming API when we eventually generate our client SDK. This will allow us to consume this API in a streaming manner whilst maintaining type-safety!
  • We use the x-query setting to indicate that while this is a POST request, we will treat it as a query instead of a mutation, allowing us to take full advantage of TanStack Query managing our streaming state.
  • Our API simply returns a stream of text as defined by both the media_type="text/plain" and the response_class=PlainTextResponse

Infrastructure

The Infrastructure we set up previously assumes that all APIs have an API Gateway integrating with a Lambda. For our story_api we actually don’t want to use API Gateway as this does not support streaming repsonses. Instead, we will use a Lambda Function URL configured with response streaming.

To support this, we are going to first update our CDK constructs as follows:

packages/common/constructs/src/core/http-api.ts
import { Construct } from 'constructs';
import { CfnOutput, Duration, Stack } from 'aws-cdk-lib';
import {
  CorsHttpMethod,
  HttpApi as _HttpApi,
  HttpMethod,
  IHttpRouteAuthorizer,
} from 'aws-cdk-lib/aws-apigatewayv2';
import { HttpLambdaIntegration } from 'aws-cdk-lib/aws-apigatewayv2-integrations';
import {
  Code,
  Function,
  FunctionUrl,
  FunctionUrlAuthType,
  InvokeMode,
  LayerVersion,
  Runtime,
  Tracing,
} from 'aws-cdk-lib/aws-lambda';
import { Grant, IGrantable } from 'aws-cdk-lib/aws-iam';
import { RuntimeConfig } from './runtime-config.js';


export interface HttpApiProps {
  readonly apiName: string;
  readonly handler: string;
  readonly handlerFilePath: string;
  readonly runtime: Runtime;
  readonly defaultAuthorizer: IHttpRouteAuthorizer;
  readonly apiType?: 'api-gateway' | 'function-url-streaming';
  readonly allowedOrigins?: string[];
}


export class HttpApi extends Construct {
  public readonly api?: _HttpApi;
  public readonly routerFunctionUrl?: FunctionUrl;
  public readonly routerFunction: Function;


  constructor(scope: Construct, id: string, props: HttpApiProps) {
    super(scope, id);


    this.routerFunction = new Function(this, `${id}Handler`, {
      timeout: Duration.seconds(30),
      runtime: props.runtime,
      handler: props.handler,
      code: Code.fromAsset(props.handlerFilePath),
      tracing: Tracing.ACTIVE,
      environment: {
        AWS_CONNECTION_REUSE_ENABLED: '1',
      },
    });


    let apiUrl;
    if (props.apiType === 'function-url-streaming') {
      const stack = Stack.of(this);
      this.routerFunction.addLayers(
        LayerVersion.fromLayerVersionArn(
          this,
          'LWALayer',
          `arn:aws:lambda:${stack.region}:753240598075:layer:LambdaAdapterLayerX86:24`,
        ),
      );
      this.routerFunction.addEnvironment('PORT', '8000');
      this.routerFunction.addEnvironment(
        'AWS_LWA_INVOKE_MODE',
        'response_stream',
      );
      this.routerFunction.addEnvironment(
        'AWS_LAMBDA_EXEC_WRAPPER',
        '/opt/bootstrap',
      );
      this.routerFunctionUrl = this.routerFunction.addFunctionUrl({
        authType: FunctionUrlAuthType.AWS_IAM,
        invokeMode: InvokeMode.RESPONSE_STREAM,
        cors: {
          allowedOrigins: props.allowedOrigins ?? ['*'],
          allowedHeaders: [
            'authorization',
            'content-type',
            'x-amz-content-sha256',
            'x-amz-date',
            'x-amz-security-token',
          ],
        },
      });
      apiUrl = this.routerFunctionUrl.url;
    } else {
      this.api = new _HttpApi(this, id, {
        corsPreflight: {
          allowOrigins: props.allowedOrigins ?? ['*'],
          allowMethods: [CorsHttpMethod.ANY],
          allowHeaders: [
            'authorization',
            'content-type',
            'x-amz-content-sha256',
            'x-amz-date',
            'x-amz-security-token',
          ],
        },
        defaultAuthorizer: props.defaultAuthorizer,
      });


      this.api.addRoutes({
        path: '/{proxy+}',
        methods: [
          HttpMethod.GET,
          HttpMethod.DELETE,
          HttpMethod.POST,
          HttpMethod.PUT,
          HttpMethod.PATCH,
          HttpMethod.HEAD,
        ],
        integration: new HttpLambdaIntegration(
          'RouterIntegration',
          this.routerFunction,
        ),
      });
      apiUrl = this.api.url;
    }


    new CfnOutput(this, `${props.apiName}Url`, { value: apiUrl! });


    RuntimeConfig.ensure(this).config.httpApis = {
      ...RuntimeConfig.ensure(this).config.httpApis!,
      [props.apiName]: apiUrl,
    };
  }


  public grantInvokeAccess(grantee: IGrantable) {
    if (this.api) {
      Grant.addToPrincipal({
        grantee,
        actions: ['execute-api:Invoke'],
        resourceArns: [this.api.arnForExecuteApi('*', '/*', '*')],
      });
    } else if (this.routerFunction) {
      Grant.addToPrincipal({
        grantee,
        actions: ['lambda:InvokeFunctionUrl'],
        resourceArns: [this.routerFunction.functionArn],
        conditions: {
          StringEquals: {
            'lambda:FunctionUrlAuthType': 'AWS_IAM',
          },
        },
      });
    }
  }
}

Now we will update the story_api to support the Lambda Web Adapter deployment.

packages/story_api/run.sh
#!/bin/bash


PATH=$PATH:$LAMBDA_TASK_ROOT/bin \
    PYTHONPATH=$PYTHONPATH:/opt/python:$LAMBDA_RUNTIME_DIR \
    exec python -m uvicorn --port=$PORT story_api.main:app

Deployment and testing

First, lets build the codebase:

Terminal window
pnpm nx run-many --target build --all

Your application can now be deployed by running the following command:

Terminal window
pnpm nx run @dungeon-adventure/infra:deploy dungeon-adventure-infra-sandbox

This deployment will take around 2 minutes to complete.

You can also deploy all stacks at once. Click here for more details.

Once the deployment completes, you should see some outputs similar to the following (some values have been redacted):

Terminal window
dungeon-adventure-infra-sandbox
dungeon-adventure-infra-sandbox: deploying... [2/2]


 ✅  dungeon-adventure-infra-sandbox


✨  Deployment time: 354s


Outputs:
dungeon-adventure-infra-sandbox.ElectroDbTableTableNameXXX = dungeon-adventure-infra-sandbox-ElectroDbTableXXX-YYY
dungeon-adventure-infra-sandbox.GameApiGameApiUrlXXX = https://xxx.region.amazonaws.com/
dungeon-adventure-infra-sandbox.GameUIDistributionDomainNameXXX = xxx.cloudfront.net
dungeon-adventure-infra-sandbox.StoryApiStoryApiUrlXXX = https://xxx.lambda-url.ap-southeast-2.on.aws/
dungeon-adventure-infra-sandbox.UserIdentityUserIdentityIdentityPoolIdXXX = region:xxx
dungeon-adventure-infra-sandbox.UserIdentityUserIdentityUserPoolIdXXX = region_xxx

We can test our API by either:

  • Starting a local instance of the FastApi server and invoke the API’s using curl.
  • Calling the deployed API using sigv4 enabled curl directly

Start your local FastAPI server by running the following command:

Terminal window
pnpm nx run dungeon_adventure.story_api:serve

Once the FastAPI server is up and running, call it by running the following command:

Terminal window
curl -N -X POST http://127.0.0.1:8000/story/generate \
  -d '{"genre":"superhero", "actions":[], "playerName":"UnnamedHero"}' \
  -H "Content-Type: application/json"

If the command executes successfully, you should see a response being streamed similar to:

UnnamedHero stood tall, his cape billowing in the wind....

Congratulations. You have built and deployed your first API using FastAPI! 🎉🎉🎉