306 lines
10 KiB
Python
306 lines
10 KiB
Python
import os
|
|
import json
|
|
import asyncio
|
|
import logging
|
|
|
|
import pymysql
|
|
import boto3
|
|
from dotenv import load_dotenv
|
|
|
|
from opensearchpy import OpenSearch
|
|
from livekit.agents import AutoSubscribe, JobContext, JobProcess, WorkerOptions, cli, llm
|
|
from livekit.agents.voice import Agent, AgentSession
|
|
from livekit.plugins import aws, silero
|
|
|
|
from engine_llm import (
|
|
EngineLLM,
|
|
engine_client_from_env,
|
|
on_playback_interrupted,
|
|
parse_metadata,
|
|
sanitize_for_tts,
|
|
)
|
|
|
|
load_dotenv()
|
|
logger = logging.getLogger("voice-agent")
|
|
|
|
VECTOR_FIELD = os.getenv("OPENSEARCH_VECTOR_FIELD", "vector_field")
|
|
TOP_K = 5
|
|
# NMSLIB cannot apply a filter inside kNN. Pull a wider set, then keep this booklet.
|
|
CANDIDATE_K = 50
|
|
|
|
os_client = OpenSearch(
|
|
hosts=[os.getenv("OPENSEARCH_URL")],
|
|
http_auth=(os.getenv("OPENSEARCH_USER"), os.getenv("OPENSEARCH_PASSWORD")),
|
|
use_ssl=True,
|
|
verify_certs=False,
|
|
ssl_show_warn=False,
|
|
)
|
|
|
|
bedrock_runtime = boto3.client("bedrock-runtime", region_name=os.getenv("AWS_REGION", "ap-south-1"))
|
|
|
|
|
|
def get_booklet_id_from_code(code: str):
|
|
connection = None
|
|
try:
|
|
connection = pymysql.connect(
|
|
host=os.getenv("DB_HOST"),
|
|
user=os.getenv("DB_USERNAME"),
|
|
password=os.getenv("DB_PASSWORD"),
|
|
database=os.getenv("DB_DATABASE"),
|
|
port=int(os.getenv("DB_PORT", 3306)),
|
|
cursorclass=pymysql.cursors.DictCursor,
|
|
)
|
|
with connection.cursor() as cursor:
|
|
cursor.execute("SELECT id, title FROM booklets WHERE code = %s", (code,))
|
|
result = cursor.fetchone()
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"Database error: {e}")
|
|
return None
|
|
finally:
|
|
try:
|
|
connection.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def embed_text(text: str):
|
|
response = bedrock_runtime.invoke_model(
|
|
modelId="amazon.titan-embed-text-v2:0",
|
|
contentType="application/json",
|
|
accept="application/json",
|
|
body=json.dumps({
|
|
"inputText": text,
|
|
"dimensions": 1024,
|
|
"normalize": True,
|
|
}),
|
|
)
|
|
return json.loads(response["body"].read())["embedding"]
|
|
|
|
|
|
def search_booklet(booklet_id: str, booklet_title: str, search_query: str) -> str:
|
|
logger.info("Searching booklet %s for: %s", booklet_id, search_query)
|
|
embedding = embed_text(search_query)
|
|
query_body = {
|
|
"size": CANDIDATE_K,
|
|
"query": {
|
|
"knn": {
|
|
VECTOR_FIELD: {
|
|
"vector": embedding,
|
|
"k": CANDIDATE_K,
|
|
}
|
|
}
|
|
},
|
|
"_source": ["text", "booklet_id", "page"],
|
|
}
|
|
res = os_client.search(index="talkingbooklet", body=query_body)
|
|
booklet_key = str(booklet_id)
|
|
hits = [
|
|
hit for hit in res["hits"]["hits"]
|
|
if str((hit.get("_source") or {}).get("booklet_id")) == booklet_key
|
|
][:TOP_K]
|
|
logger.info("OpenSearch returned %s hits for booklet %s", len(hits), booklet_id)
|
|
for hit in hits:
|
|
snippet = (hit.get("_source") or {}).get("text", "")[:80].replace("\n", " ")
|
|
logger.info("Hit score=%s snippet=%s", hit.get("_score"), snippet)
|
|
|
|
if not hits:
|
|
return "NO_MATCH"
|
|
|
|
parts = []
|
|
for hit in hits:
|
|
source = hit.get("_source") or {}
|
|
parts.append(source.get("text", ""))
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
class BookletKnowledge:
|
|
def __init__(self, booklet_id: str, booklet_title: str):
|
|
self.booklet_id = booklet_id
|
|
self.booklet_title = booklet_title
|
|
|
|
def retrieve(self, search_query: str) -> str:
|
|
try:
|
|
return search_booklet(self.booklet_id, self.booklet_title, search_query)
|
|
except Exception as e:
|
|
logger.error(f"Error searching OpenSearch: {e}")
|
|
return "NO_MATCH"
|
|
|
|
@llm.function_tool(
|
|
description="Look up more of your own content when the notes already provided do not answer the question."
|
|
)
|
|
async def search_document(self, search_query: str) -> str:
|
|
return await asyncio.to_thread(self.retrieve, search_query)
|
|
|
|
|
|
class BookletVoiceAgent(Agent):
|
|
def __init__(self, knowledge: BookletKnowledge):
|
|
self.knowledge = knowledge
|
|
instructions = (
|
|
f"You are the booklet '{knowledge.booklet_title}', speaking directly to the person in front of you. "
|
|
"Answer only the question they just asked. Do not add related facts, summaries, offers to continue, or closing remarks. "
|
|
"Speak in the first person as this booklet. Never say document, excerpt, context, passage, source, or according to. "
|
|
"Use only the private notes supplied for this turn. If those notes do not contain the answer, say you do not have that information. "
|
|
"Call search_document only when the notes are missing or about a different topic. "
|
|
"Keep the spoken answer to one or two short sentences."
|
|
)
|
|
super().__init__(instructions=instructions, tools=[knowledge.search_document])
|
|
|
|
async def on_user_turn_completed(self, turn_ctx: llm.ChatContext, new_message: llm.ChatMessage) -> None:
|
|
query = (new_message.text_content or "").strip()
|
|
if not query:
|
|
return
|
|
context = await asyncio.to_thread(self.knowledge.retrieve, query)
|
|
thin = context == "NO_MATCH"
|
|
guidance = (
|
|
"These notes do not answer the question. You may call search_document once with a shorter query. "
|
|
"If that still fails, say only that you do not have that information."
|
|
if thin
|
|
else "Answer the user's question from these private notes and nothing else. Do not mention the notes."
|
|
)
|
|
turn_ctx.add_message(
|
|
role="system",
|
|
content=f"{guidance}\n\n{context}" if not thin else guidance,
|
|
)
|
|
|
|
|
|
async def entrypoint(ctx: JobContext):
|
|
room_name = ctx.room.name
|
|
raw_meta = ctx.room.metadata or getattr(getattr(ctx, "job", None), "metadata", "") or ""
|
|
meta = parse_metadata(raw_meta)
|
|
if not meta:
|
|
logger.error("room %s has no session metadata", room_name)
|
|
return
|
|
await start_engine_voice(ctx, meta)
|
|
|
|
|
|
async def start_engine_voice(ctx: JobContext, meta: dict):
|
|
session_id = meta["sessionId"]
|
|
client = engine_client_from_env()
|
|
logger.info("sessionId=%s joining engine voice room", session_id)
|
|
await ctx.connect(auto_subscribe=AutoSubscribe.AUDIO_ONLY)
|
|
region = os.getenv("AWS_REGION", "ap-south-1")
|
|
stt_kwargs = {"region": region, "language": meta.get("language") or "en-US"}
|
|
if meta.get("vocabularyName"):
|
|
stt_kwargs["vocabulary_name"] = meta["vocabularyName"]
|
|
engine = EngineLLM(
|
|
client,
|
|
session_id,
|
|
filler_after_ms=int(os.getenv("FILLER_AFTER_MS", "800")),
|
|
)
|
|
agent = Agent(instructions="Speak the book engine's words. Do not add facts of your own.")
|
|
session = AgentSession(
|
|
vad=ctx.proc.userdata.get("vad") or silero.VAD.load(),
|
|
stt=aws.STT(**stt_kwargs),
|
|
llm=engine,
|
|
tts=aws.TTS(
|
|
voice=meta.get("voiceId") or "Joanna",
|
|
speech_engine="neural",
|
|
region=region,
|
|
),
|
|
min_endpointing_delay=float(os.getenv("MIN_ENDPOINTING_DELAY", "0.5")),
|
|
min_interruption_words=int(os.getenv("MIN_INTERRUPTION_WORDS", "2")),
|
|
user_away_timeout=float(os.getenv("USER_AWAY_TIMEOUT", "20")),
|
|
)
|
|
prompted = {"done": False}
|
|
|
|
@session.on("user_state_changed")
|
|
def on_user_state(event):
|
|
if event.new_state != "away" or prompted["done"]:
|
|
return
|
|
prompted["done"] = True
|
|
|
|
async def nudge():
|
|
await session.say("Are you still there?", allow_interruptions=True)
|
|
await asyncio.sleep(float(os.getenv("USER_AWAY_TIMEOUT", "20")))
|
|
await session.say("I'll let you go. Goodbye.", allow_interruptions=False)
|
|
try:
|
|
client.end_session(session_id)
|
|
finally:
|
|
await ctx.room.disconnect()
|
|
|
|
asyncio.create_task(nudge())
|
|
|
|
@session.on("speech_created")
|
|
def on_speech(event):
|
|
handle = event.speech_handle
|
|
|
|
def finished(done):
|
|
if not handle.interrupted:
|
|
return
|
|
on_playback_interrupted(client, session_id, engine.last_text, None)
|
|
logger.info("sessionId=%s interruption=1", session_id)
|
|
|
|
handle.add_done_callback(finished)
|
|
|
|
async def shutdown():
|
|
try:
|
|
client.end_session(session_id)
|
|
except Exception as error:
|
|
logger.error("sessionId=%s end failed: %s", session_id, error)
|
|
client.close()
|
|
|
|
ctx.add_shutdown_callback(shutdown)
|
|
await session.start(agent, room=ctx.room)
|
|
opening = sanitize_for_tts(meta.get("openingLine") or "Hello. What would you like to know?")
|
|
await session.say(opening, allow_interruptions=True)
|
|
|
|
|
|
async def start_legacy_voice(ctx: JobContext, room_name: str):
|
|
if not room_name.startswith("room_"):
|
|
logger.error(f"Invalid room name format: {room_name}")
|
|
return
|
|
|
|
code = room_name.split("_", 1)[1]
|
|
booklet_info = await asyncio.to_thread(get_booklet_id_from_code, code)
|
|
|
|
if not booklet_info:
|
|
logger.error(f"Could not find booklet with code {code}")
|
|
return
|
|
|
|
booklet_id = booklet_info["id"]
|
|
booklet_title = booklet_info["title"]
|
|
|
|
logger.info(f"Connecting to room {room_name} for booklet '{booklet_title}' (ID: {booklet_id})")
|
|
|
|
await ctx.connect(auto_subscribe=AutoSubscribe.AUDIO_ONLY)
|
|
|
|
knowledge = BookletKnowledge(str(booklet_id), booklet_title)
|
|
agent = BookletVoiceAgent(knowledge)
|
|
|
|
session = AgentSession(
|
|
vad=ctx.proc.userdata.get("vad") or silero.VAD.load(),
|
|
stt=aws.STT(region=os.getenv("AWS_REGION", "ap-south-1")),
|
|
llm=aws.LLM(
|
|
model="anthropic.claude-3-haiku-20240307-v1:0",
|
|
region=os.getenv("AWS_REGION", "ap-south-1"),
|
|
),
|
|
tts=aws.TTS(
|
|
voice="Joanna",
|
|
speech_engine="neural",
|
|
region=os.getenv("AWS_REGION", "ap-south-1"),
|
|
),
|
|
)
|
|
|
|
await session.start(agent, room=ctx.room)
|
|
|
|
await asyncio.sleep(1)
|
|
await session.say(
|
|
"Hello. What would you like to know?",
|
|
allow_interruptions=True,
|
|
)
|
|
|
|
|
|
def prewarm(proc: JobProcess):
|
|
proc.userdata["vad"] = silero.VAD.load()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli.run_app(
|
|
WorkerOptions(
|
|
entrypoint_fnc=entrypoint,
|
|
prewarm_fnc=prewarm,
|
|
agent_name=os.getenv("LIVEKIT_AGENT_NAME", ""),
|
|
)
|
|
)
|