livekit--agents
wehub-resource-sync
107 行
3.4 KiB
Python
107 行
3.4 KiB
Python
from pathlib import Path
|
|
|
|
from dotenv import load_dotenv
|
|
from llama_index.core import (
|
|
SimpleDirectoryReader,
|
|
StorageContext,
|
|
VectorStoreIndex,
|
|
load_index_from_storage,
|
|
)
|
|
from llama_index.core.schema import MetadataMode
|
|
|
|
from livekit.agents import (
|
|
Agent,
|
|
AgentServer,
|
|
AgentSession,
|
|
AutoSubscribe,
|
|
JobContext,
|
|
cli,
|
|
llm,
|
|
)
|
|
from livekit.agents.voice.agent import ModelSettings
|
|
from livekit.plugins import deepgram, openai
|
|
|
|
load_dotenv()
|
|
|
|
# check if storage already exists
|
|
THIS_DIR = Path(__file__).parent
|
|
PERSIST_DIR = THIS_DIR / "retrieval-engine-storage"
|
|
if not PERSIST_DIR.exists():
|
|
# load the documents and create the index
|
|
documents = SimpleDirectoryReader(THIS_DIR / "data").load_data()
|
|
index = VectorStoreIndex.from_documents(documents)
|
|
# store it for later
|
|
index.storage_context.persist(persist_dir=PERSIST_DIR)
|
|
else:
|
|
# load the existing index
|
|
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
|
|
index = load_index_from_storage(storage_context)
|
|
|
|
|
|
class RetrievalAgent(Agent):
|
|
def __init__(self, index: VectorStoreIndex):
|
|
super().__init__(
|
|
instructions=(
|
|
"You are a voice assistant created by LiveKit. Your interface "
|
|
"with users will be voice. You should use short and concise "
|
|
"responses, and avoiding usage of unpronouncable punctuation."
|
|
),
|
|
stt=deepgram.STT(),
|
|
llm=openai.LLM(),
|
|
tts=openai.TTS(),
|
|
)
|
|
self.index = index
|
|
|
|
async def llm_node(
|
|
self,
|
|
chat_ctx: llm.ChatContext,
|
|
tools: list[llm.FunctionTool],
|
|
model_settings: ModelSettings,
|
|
):
|
|
user_msg = chat_ctx.items[-1]
|
|
assert isinstance(user_msg, llm.ChatMessage) and user_msg.role == "user"
|
|
user_query = user_msg.text_content
|
|
assert user_query is not None
|
|
|
|
retriever = self.index.as_retriever()
|
|
nodes = await retriever.aretrieve(user_query)
|
|
|
|
instructions = "Context that might help answer the user's question:"
|
|
for node in nodes:
|
|
node_content = node.get_content(metadata_mode=MetadataMode.LLM)
|
|
instructions += f"\n\n{node_content}"
|
|
|
|
# update the instructions for this turn, you may use some different methods
|
|
# to inject the context into the chat_ctx that fits the LLM you are using
|
|
system_msg = chat_ctx.items[0]
|
|
if isinstance(system_msg, llm.ChatMessage) and system_msg.role == "system":
|
|
# TODO(long): provide an api to update the instructions of chat_ctx
|
|
system_msg.content.append(instructions)
|
|
else:
|
|
chat_ctx.items.insert(0, llm.ChatMessage(role="system", content=[instructions]))
|
|
preview = instructions[:100].replace("\n", "\\n")
|
|
print(f"update instructions: {preview}...")
|
|
|
|
# update the instructions for agent
|
|
# await self.update_instructions(instructions)
|
|
|
|
return Agent.default.llm_node(self, chat_ctx, tools, model_settings)
|
|
|
|
|
|
server = AgentServer()
|
|
|
|
|
|
@server.rtc_session()
|
|
async def entrypoint(ctx: JobContext):
|
|
await ctx.connect(auto_subscribe=AutoSubscribe.AUDIO_ONLY)
|
|
|
|
agent = RetrievalAgent(index)
|
|
session = AgentSession()
|
|
await session.start(agent=agent, room=ctx.room)
|
|
|
|
await session.say("Hey, how can I help you today?", allow_interruptions=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli.run_app(server)
|