Github/screenshot-to-code
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

AWS Bedrock support - adjust code style

Browse Source
This commit is contained in:
yoreland 2024-06-09 00:01:32 +08:00
parent 7032342f91
commit afd34dfaad
4 changed files with 179 additions and 184 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
backend/config.py
View File

@ -4,8 +4,11 @@
import os
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY", None)
# AWS
AWS_ACCESS_KEY = os.environ.get("AWS_ACCESS_KEY", None)
AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY", None)
AWS_REGION_NAME = os.environ.get("AWS_REGION_NAME", "us-west-2")
# Debugging-related

4
backend/evals/core.py
View File

@ -1,4 +1,6 @@
import os
from config import AWS_REGION_NAME
from config import AWS_ACCESS_KEY
from config import AWS_SECRET_ACCESS_KEY
from config import ANTHROPIC_API_KEY
@ -14,6 +16,7 @@ async def generate_code_core(image_url: str, stack: Stack, model: Llm) -> str:
anthropic_api_key = ANTHROPIC_API_KEY
aws_access_key = AWS_ACCESS_KEY
aws_secret_access_key = AWS_SECRET_ACCESS_KEY
aws_region_name = AWS_REGION_NAME,
openai_base_url = None
async def process_chunk(content: str):
@ -34,6 +37,7 @@ async def generate_code_core(image_url: str, stack: Stack, model: Llm) -> str:
prompt_messages,
access_key=aws_access_key,
secret_access_key=aws_secret_access_key,
aws_region_name=aws_region_name,
callback=lambda x: process_chunk(x),
)
else:

199
backend/llm.py
View File

@ -17,6 +17,7 @@ class Llm(Enum):
GPT_4_VISION = "gpt-4-vision-preview"
GPT_4_TURBO_2024_04_09 = "gpt-4-turbo-2024-04-09"
CLAUDE_3_SONNET = "claude-3-sonnet-20240229"
CLAUDE_3_SONNET_BEDROCK = "anthropic.claude-3-sonnet-20240229-v1:0"
CLAUDE_3_OPUS = "claude-3-opus-20240229"
CLAUDE_3_HAIKU = "claude-3-haiku-20240307"
@ -125,61 +126,42 @@ async def stream_claude_response(
return response.content[0].text
async def stream_claude_response_aws_bedrock(
messages: List[ChatCompletionMessageParam],
access_key: str,
secret_access_key: str,
callback: Callable[[str], Awaitable[None]],
) -> str:
def initialize_bedrock_client(access_key: str, secret_access_key: str, aws_region_name: str):
try:
# Initialize the Bedrock Runtime client
bedrock_runtime = boto3.client(
service_name='bedrock-runtime',
aws_access_key_id=access_key,
aws_secret_access_key=secret_access_key,
region_name=aws_region_name,
)
return bedrock_runtime
except ClientError as err:
message = err.response["Error"]["Message"]
print(f"A client error occurred: {message}")
except Exception as err:
print("An error occurred!")
raise err
# Set model parameters
model_id = 'anthropic.claude-3-sonnet-20240229-v1:0'
max_tokens = 4096
content_type = 'application/json'
accept = '*/*'
# Translate OpenAI messages to Claude messages
system_prompt = cast(str, messages[0].get("content"))
claude_messages = [dict(message) for message in messages[1:]]
for message in claude_messages:
if not isinstance(message["content"], list):
continue
for content in message["content"]: # type: ignore
if content["type"] == "image_url":
content["type"] = "image"
# Extract base64 data and media type from data URL
# Example base64 data URL: data:image/png;base64,iVBOR...
image_data_url = cast(str, content["image_url"]["url"])
media_type = image_data_url.split(";")[0].split(":")[1]
base64_data = image_data_url.split(",")[1]
# Remove OpenAI parameter
del content["image_url"]
content["source"] = {
"type": "base64",
"media_type": media_type,
"data": base64_data,
}
async def stream_bedrock_response(
bedrock_runtime,
messages: List[dict],
system_prompt: str,
model_id: str,
max_tokens: int,
content_type: str,
accept: str,
temperature: float,
callback: Callable[[str], Awaitable[None]],
) -> str:
try:
# Prepare the request body
body = json.dumps({
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": max_tokens,
"messages": claude_messages,
"messages": messages,
"system": system_prompt,
"temperature": 0.0
"temperature": temperature
})
# Invoke the Bedrock Runtime API with response stream
@ -213,6 +195,60 @@ async def stream_claude_response_aws_bedrock(
print("An error occurred!")
raise err
async def stream_claude_response_aws_bedrock(
messages: List[dict],
access_key: str,
secret_access_key: str,
aws_region_name: str,
callback: Callable[[str], Awaitable[None]],
) -> str:
bedrock_runtime = initialize_bedrock_client(access_key, secret_access_key, aws_region_name)
# Set model parameters
model_id = Llm.CLAUDE_3_SONNET_BEDROCK.value
max_tokens = 4096
content_type = 'application/json'
accept = '*/*'
temperature = 0.0
# Translate OpenAI messages to Claude messages
system_prompt = cast(str, messages[0].get("content"))
claude_messages = [dict(message) for message in messages[1:]]
for message in claude_messages:
if not isinstance(message["content"], list):
continue
for content in message["content"]: # type: ignore
if content["type"] == "image_url":
content["type"] = "image"
# Extract base64 data and media type from data URL
# Example base64 data URL: data:image/png;base64,iVBOR...
image_data_url = cast(str, content["image_url"]["url"])
media_type = image_data_url.split(";")[0].split(":")[1]
base64_data = image_data_url.split(",")[1]
# Remove OpenAI parameter
del content["image_url"]
content["source"] = {
"type": "base64",
"media_type": media_type,
"data": base64_data,
}
return await stream_bedrock_response(
bedrock_runtime,
claude_messages,
system_prompt,
model_id,
max_tokens,
content_type,
accept,
temperature,
callback,
)
async def stream_claude_response_native(
system_prompt: str,
messages: list[Any],
@ -306,22 +342,15 @@ async def stream_claude_response_native_aws_bedrock(
messages: list[Any],
access_key: str,
secret_access_key: str,
aws_region_name: str,
callback: Callable[[str], Awaitable[None]],
include_thinking: bool = False,
model: Llm = Llm.CLAUDE_3_OPUS,
model: Llm = Llm.CLAUDE_3_SONNET_BEDROCK,
) -> str:
try:
# Initialize the Bedrock Runtime client
bedrock_runtime = boto3.client(
service_name='bedrock-runtime',
aws_access_key_id=access_key,
aws_secret_access_key=secret_access_key,
)
bedrock_runtime = initialize_bedrock_client(access_key, secret_access_key, aws_region_name)
# Set model parameters
# model_id = model.value
model_id = "anthropic.claude-3-sonnet-20240229-v1:0"
model_id = Llm.CLAUDE_3_SONNET_BEDROCK.value
max_tokens = 4096
content_type = 'application/json'
accept = '*/*'
@ -334,10 +363,6 @@ async def stream_claude_response_native_aws_bedrock(
prefix = "<thinking>"
response = None
# For debugging
full_stream = ""
debug_file_writer = DebugFileWriter()
while current_pass_num <= max_passes:
current_pass_num += 1
@ -348,41 +373,17 @@ async def stream_claude_response_native_aws_bedrock(
else messages
)
pprint_prompt(messages_to_send)
# Prepare the request body
body = json.dumps({
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": max_tokens,
"messages": messages_to_send,
"system": system_prompt,
"temperature": temperature
})
# Invoke the Bedrock Runtime API with response stream
response = bedrock_runtime.invoke_model_with_response_stream(
body=body,
modelId=model_id,
accept=accept,
contentType=content_type,
response_text = await stream_bedrock_response(
bedrock_runtime,
messages_to_send,
system_prompt,
model_id,
max_tokens,
content_type,
accept,
temperature,
callback,
)
stream = response.get("body")
# Stream the response
response_text = ""
if stream:
for event in stream:
chunk = event.get("chunk")
if chunk:
data = chunk.get("bytes").decode()
chunk_obj = json.loads(data)
if chunk_obj["type"] == "content_block_delta":
text = chunk_obj["delta"]["text"]
response_text += text
print(text, end="", flush=True)
full_stream += text
await callback(text)
# Set up messages array for next pass
messages += [
@ -393,18 +394,4 @@ async def stream_claude_response_native_aws_bedrock(
},
]
# print(
# f"Token usage: Input Tokens: {response.usage.input_tokens}, Output Tokens: {response.usage.output_tokens}"
# )
if not response:
raise Exception("No HTML response found in AI response")
else:
return full_stream
except ClientError as err:
message = err.response["Error"]["Message"]
print(f"A client error occurred: {message}")
except Exception as err:
print("An error occurred!")
raise err
return response_text

5
backend/routes/generate_code.py
View File

@ -2,7 +2,7 @@ import os
import traceback
from fastapi import APIRouter, WebSocket
import openai
from config import ANTHROPIC_API_KEY, IS_PROD, SHOULD_MOCK_AI_RESPONSE, AWS_ACCESS_KEY, AWS_SECRET_ACCESS_KEY
from config import ANTHROPIC_API_KEY, IS_PROD, SHOULD_MOCK_AI_RESPONSE, AWS_ACCESS_KEY, AWS_SECRET_ACCESS_KEY, AWS_REGION_NAME
from custom_types import InputMode
from llm import (
Llm,
@ -25,7 +25,6 @@ from prompts.types import Stack
from video.utils import extract_tag_content, assemble_claude_prompt_video
from ws.constants import APP_ERROR_WEB_SOCKET_CODE # type: ignore
router = APIRouter()
@ -239,6 +238,7 @@ async def stream_code(websocket: WebSocket):
messages=prompt_messages, # type: ignore
access_key=AWS_ACCESS_KEY,
secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_region_name=AWS_REGION_NAME,
callback=lambda x: process_chunk(x),
model=Llm.CLAUDE_3_OPUS,
include_thinking=True,
@ -261,6 +261,7 @@ async def stream_code(websocket: WebSocket):
prompt_messages, # type: ignore
access_key=AWS_ACCESS_KEY,
secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_region_name=AWS_REGION_NAME,
callback=lambda x: process_chunk(x),
)
exact_llm_version = code_generation_model