python-agents-examples/docs/examples/rag/main.py at main · charan632-dev/python-agents-examples · GitHub

History

257 lines (212 loc) · 8.01 KB

Raw

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

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

#!/usr/bin/env python3

"""

---

title: RAG Voice Agent

category: rag

tags: [annoy_index, vector_search, embeddings, livekit_docs, deduplication]

difficulty: advanced

description: RAG-enabled agent with vector search for LiveKit documentation

demonstrates:

- Annoy vector index loading and querying

- OpenAI embeddings for semantic search

- Result deduplication with seen tracking

- Function tool for document search

- Paragraph-based context retrieval

- Noise cancellation with BVC

---

"""

RAG-enabled voice agent example for LiveKit Agents 1.0

This agent uses the RAG (Retrieval Augmented Generation) plugin to provide

information from a knowledge base when answering user questions.

Before running this agent:

1. Make sure you have your OpenAI API key in a .env file

2. Run build_rag_data.py to build the RAG database

"""

import logging

import pickle

from pathlib import Path

from typing import Literal, Any

from collections.abc import Iterable

from dataclasses import dataclass

from dotenv import load_dotenv

import annoy

from livekit.agents import (

JobContext,

WorkerOptions,

cli,

RunContext,

function_tool,

RoomInputOptions,

Agent,

AgentSession, inference,

)

from livekit.plugins import openai, silero, deepgram, noise_cancellation

from livekit.plugins.turn_detector.english import EnglishModel

# Load environment variables

load_dotenv()

# Configure logging

logging.basicConfig(

level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"

)

logger = logging.getLogger("rag-agent")

# RAG Index Types and Classes

Metric = Literal["angular", "euclidean", "manhattan", "hamming", "dot"]

ANNOY_FILE = "index.annoy"

METADATA_FILE = "metadata.pkl"

@dataclass

class _FileData:

f: int

metric: Metric

userdata: dict[int, Any]

@dataclass

class Item:

i: int

userdata: Any

vector: list[float]

@dataclass

class QueryResult:

userdata: Any

distance: float

class AnnoyIndex:

def __init__(self, index: annoy.AnnoyIndex, filedata: _FileData) -> None:

self._index = index

self._filedata = filedata

@classmethod

def load(cls, path: str) -> "AnnoyIndex":

p = Path(path)

index_path = p / ANNOY_FILE

metadata_path = p / METADATA_FILE

with open(metadata_path, "rb") as f:

metadata: _FileData = pickle.load(f)

index = annoy.AnnoyIndex(metadata.f, metadata.metric)

index.load(str(index_path))

return cls(index, metadata)

@property

def size(self) -> int:

return self._index.get_n_items()

def items(self) -> Iterable[Item]:

for i in range(self._index.get_n_items()):

item = Item(

i=i,

userdata=self._filedata.userdata[i],

vector=self._index.get_item_vector(i),

)

yield item

def query(

self, vector: list[float], n: int, search_k: int = -1

) -> list[QueryResult]:

ids = self._index.get_nns_by_vector(

vector, n, search_k=search_k, include_distances=True

)

return [

QueryResult(userdata=self._filedata.userdata[i], distance=distance)

for i, distance in zip(*ids)

]

class RAGEnrichedAgent(Agent):

"""

An agent that can answer questions using RAG (Retrieval Augmented Generation).

"""

def __init__(self) -> None:

"""Initialize the RAG-enabled agent."""

super().__init__(

instructions="""

You are a helpful voice assistant specializing in knowledge about LiveKit ("live" pronounced as in "live stream").

You can answer questions about LiveKit, the LiveKit Agents SDK, and more.

Your responses should always be concise and suitable for text-to-speech output, so be casual and avoid using markdown or other special formatting.

""",

)

# Initialize RAG components

vdb_dir = Path(__file__).parent / "data"

data_path = vdb_dir / "paragraphs.pkl"

if not vdb_dir.exists() or not data_path.exists():

logger.warning(

"RAG database not found. Please run build_rag_data.py first:\n"

"$ python build_rag_data.py"

)

return

# Load RAG index and data

self._index_path = vdb_dir

self._data_path = data_path

self._embeddings_dimension = 1536

self._embeddings_model = "text-embedding-3-small"

self._seen_results = set() # Track previously seen results

try:

self._annoy_index = AnnoyIndex.load(str(self._index_path))

with open(self._data_path, "rb") as f:

self._paragraphs_by_uuid = pickle.load(f)

logger.info("RAG database loaded successfully.")

except Exception as e:

logger.error(f"Failed to load RAG database: {e}")

@function_tool

async def livekit_docs_search(self, context: RunContext, query: str):

"""Lookup information in the LiveKit docs database. Will not return results already returned in previous lookups."""

try:

# Generate embeddings for the query

query_embedding = await openai.create_embeddings(

input=[query],

model=self._embeddings_model,

dimensions=self._embeddings_dimension,

)

# Query the index for more results than we need to ensure we have enough new content

all_results = self._annoy_index.query(

query_embedding[0].embedding, n=5

) # Get more results initially

# Filter out previously seen results

new_results = [

r for r in all_results if r.userdata not in self._seen_results

]

# If we don't have enough new results, clear the seen results and start fresh

if len(new_results) == 0:

return "No new results found."

else:

new_results = new_results[:2] # Take top 2 new results

# Build context from multiple relevant paragraphs

context_parts = []

for result in new_results:

# Add result to seen set

self._seen_results.add(result.userdata)

paragraph = self._paragraphs_by_uuid.get(result.userdata, "")

if paragraph:

# Extract source URL if available in the paragraph

source = "Unknown source"

if "from [" in paragraph:

source = paragraph.split("from [")[1].split("]")[0]

paragraph = paragraph.split("]")[1].strip()

context_parts.append(f"Source: {source}\nContent: {paragraph}\n")

if not context_parts:

return

# Combine all context parts with clear separation

full_context = "\n\n".join(context_parts)

logger.info(

f"Results for query: {query}, full context: {full_context.replace('\n', '\\n')}"

)

return full_context

except Exception as e:

return "Could not find any relevant information for that query."

async def on_enter(self):

"""Called when the agent enters the session."""

self.session.generate_reply(

instructions="Briefly greet the user and offer your assistance with LiveKit."

)

async def entrypoint(ctx: JobContext):

"""Main entrypoint for the agent."""

session = AgentSession(

stt=deepgram.STT(),

llm=openai.LLM(model="gpt-4o"),

tts=openai.TTS(

instructions="You are a helpful assistant with a pleasant voice.",

voice="ash",

turn_detection=EnglishModel(),

vad=silero.VAD.load(),

)

await session.start(

agent=RAGEnrichedAgent(),

room=ctx.room,

room_input_options=RoomInputOptions(

noise_cancellation=noise_cancellation.BVC(),

)

if __name__ == "__main__":

cli.run_app(WorkerOptions(entrypoint_fnc=entrypoint))

python-agents-examples/docs/examples/rag/main.py at main · charan632-dev/python-agents-examples · GitHub

Expand file tree

Breadcrumbs

main.py

Latest commit

History

Breadcrumbs

main.py

File metadata and controls