Spaces:
Running
on
T4
Running
on
T4
File size: 5,996 Bytes
2e2dda5 5a7796a 2e2dda5 ccd79d7 2e2dda5 3e0e0e3 2e2dda5 ccd79d7 c572109 ccd79d7 2625508 ccd79d7 2e2dda5 3e0e0e3 2e2dda5 3e0e0e3 2e2dda5 8a64b48 2e2dda5 ccd79d7 2e2dda5 4df48e4 2e2dda5 ccd79d7 2e2dda5 ccd79d7 2e2dda5 eb03410 2e2dda5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import boto3
import json
import time
import zipfile
from io import BytesIO
import uuid
import pprint
import logging
from PIL import Image
import os
import base64
import re
import requests
#import utilities.re_ranker as re_ranker
import utilities.invoke_models as invoke_models
import streamlit as st
import time as t
import botocore.exceptions
from datetime import datetime, timezone
import botocore
import utilities.ubi_lambda as ubi
if "inputs_" not in st.session_state:
st.session_state.inputs_ = {}
parent_dirname = "/".join((os.path.dirname(__file__)).split("/")[0:-1])
region = 'us-east-1'
# setting logger
logging.basicConfig(format='[%(asctime)s] p%(process)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)
# getting boto3 clients for required AWS services
#bedrock_agent_client = boto3.client('bedrock-agent',region_name=region)
bedrock_agent_runtime_client = boto3.client(
'bedrock-agent-runtime',
aws_access_key_id=st.secrets['user_access_key_us_west_2'],
aws_secret_access_key=st.secrets['user_secret_key_us_west_2'], region_name = 'us-west-2'
)
enable_trace:bool = True
end_session:bool = False
def now_rfc3339():
return datetime.now(timezone.utc).isoformat()
def send_otel_span(span):
try:
status = ubi.send_to_lambda("otel-v1-apm-span-default", span)
if status == 202:
print("Traces sent to Lambda")
else:
print("Lambda did not accept the request")
print(f"[OTEL SPAN] {span['name']} -> {status}")
except Exception as e:
print(f"[OTEL ERROR] {e}")
def convert_to_span(block, trace_id, index):
span_id = str(uuid.uuid4()).replace("-", "")[:16]
name = "step"
attributes = {}
if "invocationInput" in block:
name = block["invocationInput"].get("function", "invocation")
attributes = {p["name"]: p["value"] for p in block["invocationInput"].get("parameters", [])}
elif "observation" in block:
name = block["observation"].get("type", "observation").lower()
attributes = block["observation"].get("actionGroupInvocationOutput", {})
elif "thinking" in block:
name = "thinking"
attributes["message"] = block["thinking"].get("content", "")
elif "rationale" in block:
name = "rationale"
attributes["message"] = block["rationale"]
return {
"traceId": trace_id,
"spanId": span_id,
"name": name,
"startTime": now_rfc3339(),
"endTime": now_rfc3339(),
"durationInNanos": 10000000 * (index + 1),
"kind": "INTERNAL",
"status": {"code": "OK"},
"attributes": attributes,
"resource": {
"service.name": "bedrock-agent"
}
}
def delete_memory():
response = bedrock_agent_runtime_client.delete_agent_memory(
agentAliasId='DEEEEZM2TM',
agentId='EJVGQW1BH7'
)
def query_(inputs):
# invoke the agent API
agentResponse = bedrock_agent_runtime_client.invoke_agent(
inputText=inputs['shopping_query'],
agentId='EJVGQW1BH7',
agentAliasId='DEEEEZM2TM',
sessionId=st.session_state.session_id_,
enableTrace=enable_trace,
endSession= end_session
)
logger.info(pprint.pprint(agentResponse))
#print("***agent*****response*********")
#print(agentResponse)
event_stream = agentResponse['completion']
total_context = []
last_tool = ""
last_tool_name = ""
agent_answer = ""
trace_id = str(uuid.uuid4()).replace("-", "")
try:
for i,event in enumerate(event_stream):
if 'trace' in event:
if('orchestrationTrace' not in event['trace']['trace']):
continue
orchestration_trace = event['trace']['trace']['orchestrationTrace']
total_context_item = {}
if('modelInvocationOutput' in orchestration_trace and '<tool_name>' in orchestration_trace['modelInvocationOutput']['rawResponse']['content']):
total_context_item['tool'] = orchestration_trace['modelInvocationOutput']['rawResponse']
if('rationale' in orchestration_trace):
total_context_item['rationale'] = orchestration_trace['rationale']['text']
if('invocationInput' in orchestration_trace):
total_context_item['invocationInput'] = orchestration_trace['invocationInput']['actionGroupInvocationInput']
last_tool_name = total_context_item['invocationInput']['function']
if('observation' in orchestration_trace):
total_context_item['observation'] = event['trace']['trace']['orchestrationTrace']['observation']
tool_output_last_obs = event['trace']['trace']['orchestrationTrace']['observation']
if(tool_output_last_obs['type'] == 'ACTION_GROUP'):
last_tool = tool_output_last_obs['actionGroupInvocationOutput']['text']
if(tool_output_last_obs['type'] == 'FINISH'):
agent_answer = tool_output_last_obs['finalResponse']['text']
if('modelInvocationOutput' in orchestration_trace and '<thinking>' in orchestration_trace['modelInvocationOutput']['rawResponse']['content']):
total_context_item['thinking'] = orchestration_trace['modelInvocationOutput']['rawResponse']
if(total_context_item!={}):
total_context.append(total_context_item)
# 🔁 Generate + send OpenTelemetry span for each block
span = convert_to_span(total_context_item, trace_id, i)
send_otel_span(span)
except botocore.exceptions.EventStreamError as error:
raise error
return {'text':agent_answer,'source':total_context,'last_tool':{'name':last_tool_name,'response':last_tool}}
|