update bot logic
This commit is contained in:
parent
aa510e99cc
commit
62aff564da
@ -1,113 +1,101 @@
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
import boto3
|
||||
import os
|
||||
import asyncio
|
||||
from fastapi.responses import StreamingResponse
|
||||
from botocore.exceptions import ClientError
|
||||
const WickrIOBotAPI = require('wickrio-bot-api');
|
||||
const WickrUser = WickrIOBotAPI.WickrUser;
|
||||
const bot = new WickrIOBotAPI.WickrIOBot();
|
||||
const WickrIOAPI = bot.apiService().WickrIOAPI;
|
||||
|
||||
import time
|
||||
const axios = require('axios');
|
||||
|
||||
app = FastAPI()
|
||||
var fs = require('fs');
|
||||
|
||||
bots_config = {
|
||||
"dronebot": {
|
||||
"agent_id": "1",
|
||||
"agent_alias": "1"
|
||||
},
|
||||
"trainerbot": {
|
||||
"agent_id": "1",
|
||||
"agent_alias": "1"
|
||||
},
|
||||
"checkinbot": {
|
||||
"agent_id": "1",
|
||||
"agent_alias": "1"
|
||||
process.stdin.resume(); //so the program will not close instantly
|
||||
|
||||
bot.processesJsonToProcessEnv()
|
||||
|
||||
var bot_username;
|
||||
|
||||
async function exitHandler(options, err) {
|
||||
try {
|
||||
var closed = await bot.close();
|
||||
if (err || options.exit) {
|
||||
console.log("Exit reason:", err);
|
||||
process.exit();
|
||||
} else if (options.pid) {
|
||||
process.kill(process.pid);
|
||||
}
|
||||
} catch (err) {
|
||||
console.log(err);
|
||||
}
|
||||
}
|
||||
ssm = boto3.client("ssm", 'us-east-1')
|
||||
|
||||
//catches ctrl+c and stop.sh events
|
||||
process.on('SIGINT', exitHandler.bind(null, {exit: true}));
|
||||
|
||||
// catches "kill pid" (for example: nodemon restart)
|
||||
process.on('SIGUSR1', exitHandler.bind(null, {pid: true}));
|
||||
process.on('SIGUSR2', exitHandler.bind(null, {pid: true}));
|
||||
|
||||
//catches uncaught exceptions
|
||||
process.on('uncaughtException', exitHandler.bind(null, {exit: true}));
|
||||
|
||||
async function main() {
|
||||
try {
|
||||
var tokens = JSON.parse(process.env.tokens);
|
||||
var status = await bot.start(tokens.WICKRIO_BOT_NAME.value)
|
||||
if (!status) {
|
||||
exitHandler(null, {
|
||||
exit: true,
|
||||
reason: 'Client not able to start'
|
||||
});
|
||||
}
|
||||
|
||||
//Passes a callback function that will receive incoming messages into the bot client
|
||||
bot.startListening(listen);
|
||||
|
||||
} catch (err) {
|
||||
return console.log(err);
|
||||
}
|
||||
}
|
||||
|
||||
async function listen(message) {
|
||||
console.log('---------------------');
|
||||
var wickrUser;
|
||||
|
||||
var parsedMessage = bot.parseMessage(message);
|
||||
if (!parsedMessage) {
|
||||
return;
|
||||
}
|
||||
|
||||
var vGroupID = parsedMessage.vgroupid;
|
||||
var userEmail = parsedMessage.userEmail;
|
||||
var convoType = parsedMessage.convoType;
|
||||
var personal_vGroupID = "";
|
||||
console.log(userEmail);
|
||||
console.log(parsedMessage.message)
|
||||
|
||||
if (convoType === 'personal')
|
||||
personal_vGroupID = vGroupID;
|
||||
|
||||
try {
|
||||
const agent_payload = {
|
||||
question: parsedMessage.message,
|
||||
requestSessionId: vGroupID,
|
||||
bot: "trainerbot"
|
||||
};
|
||||
|
||||
const agent_gateway_response = await axios.post(
|
||||
'http://172.17.0.1:8000/bedrock-agent',
|
||||
agent_payload
|
||||
);
|
||||
|
||||
await WickrIOAPI.cmdSendRoomMessage(vGroupID, agent_gateway_response.data);
|
||||
} catch (err) {
|
||||
console.error("Error:", err.message);
|
||||
await WickrIOAPI.cmdSendRoomMessage(vGroupID, err.message);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
async def fetch_agent_parameters():
|
||||
while True:
|
||||
for i in ['trainerbot', 'dronebot', 'checkinbot']:
|
||||
bot_key = i
|
||||
try:
|
||||
alias_param = f"/{bot_key}/agentalias"
|
||||
response = ssm.get_parameter(Name=alias_param, WithDecryption=True)
|
||||
bots_config[bot_key]["agent_alias"] = response['Parameter']['Value']
|
||||
except ClientError as e:
|
||||
print(f"Error fetching {alias_param}: {e}")
|
||||
bots_config[bot_key]["agent_alias"] = "1" # fallback
|
||||
|
||||
try:
|
||||
id_param = f"/{bot_key}/agentid"
|
||||
response = ssm.get_parameter(Name=id_param, WithDecryption=True)
|
||||
bots_config[bot_key]["agent_id"] = response['Parameter']['Value']
|
||||
except ClientError as e:
|
||||
print(f"Error fetching {id_param}: {e}")
|
||||
bots_config[bot_key]["agent_id"] = "1" # fallback
|
||||
|
||||
print("Updated from SSM —")
|
||||
for bot_key, config in bots_config.items():
|
||||
print(f"{bot_key.upper()} → ID: {config['agent_id']}, Alias: {config['agent_alias']}")
|
||||
|
||||
await asyncio.sleep(60)
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
asyncio.create_task(fetch_agent_parameters())
|
||||
|
||||
# Define request body
|
||||
class BedrockRequest(BaseModel):
|
||||
question: str
|
||||
requestSessionId: str = 'user'
|
||||
bot: str = 'dronebot'
|
||||
|
||||
# AWS Bedrock client setup
|
||||
def get_bedrock_client():
|
||||
return boto3.client(
|
||||
'bedrock-agent-runtime',
|
||||
region_name='us-east-1'
|
||||
)
|
||||
|
||||
@app.get("/health")
|
||||
def health_check():
|
||||
return {"message": "OK"}
|
||||
|
||||
@app.post("/bedrock-agent")
|
||||
def call_bedrock_agent(payload: BedrockRequest):
|
||||
bot_alias = bots_config[payload.bot]['agent_alias']
|
||||
bot_id = bots_config[payload.bot]['agent_id']
|
||||
client = get_bedrock_client()
|
||||
print(payload)
|
||||
try:
|
||||
response_stream = client.invoke_agent(
|
||||
sessionId=payload.requestSessionId + bot_id + bot_alias,
|
||||
agentId=bot_id,
|
||||
agentAliasId = bot_alias,
|
||||
inputText=payload.question,
|
||||
)
|
||||
|
||||
def event_stream():
|
||||
for event in response_stream["completion"]:
|
||||
if "chunk" in event:
|
||||
# Decode bytes to string if needed
|
||||
chunk_bytes = event["chunk"].get("bytes")
|
||||
if chunk_bytes:
|
||||
yield chunk_bytes.decode("utf-8")
|
||||
else:
|
||||
yield event["chunk"].get("text", "")
|
||||
|
||||
return StreamingResponse(event_stream(), media_type="text/plain")
|
||||
except Exception as e:
|
||||
print(e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get("/test-stream")
|
||||
def test_stream():
|
||||
def generator():
|
||||
for i in range(5):
|
||||
yield f"Chunk {i}\n"
|
||||
time.sleep(1)
|
||||
return StreamingResponse(generator(), media_type="text/plain")
|
||||
main();
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user