音频也能“对话”？用 AssemblyAI、Qdrant 和 DeepSeek-R1 构建音频 RAG 聊天机器人 原创 精华

git clone https://github.com/karthikponna/chat_with_audios.git
cd chat_with_audios

# 创建虚拟环境
python3 -m venv venv
source venv/bin/activate  # macOS 和 Linux
# 对于 Windows：
# python -m venv venv
# .\venv\Scripts\activate

# 安装依赖
pip install -r requirements.txt

# 设置环境变量
touch .env  # macOS 和 Linux
# 对于 Windows：
# type nul > .env
echo"ASSEMBLYAI_API_KEY=your_assemblyai_api_key_string" >> .env
echo"SAMBANOVA_API_KEY=your_sambanova_api_key_string" >> .env
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.

from typing import List
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

def batch_iterate(lst: List[str], batch_size: int) -> List[List[str]]:
    """将文本列表分割成小块"""
    for i in range(0, len(lst), batch_size):
        yield lst[i : i + batch_size]

class EmbedData:
    def __init__(self, embed_model_name: str = "BAAI/bge-large-en-v1.5", batch_size: int = 32):
        self.embed_model_name = embed_model_name
        self.embed_model = self._load_embed_model()
        self.batch_size = batch_size
        self.embeddings = []

    def _load_embed_model(self) -> HuggingFaceEmbedding:
        """加载 Hugging Face 嵌入模型"""
        embed_model = HuggingFaceEmbedding(model_name=self.embed_model_name, trust_remote_code=True, cache_folder='./hf_cache')
        return embed_model

    def generate_embedding(self, context: List[str]) -> List[List[float]]:
        """生成文本嵌入向量"""
        return self.embed_model.get_text_embedding_batch(context)

    def embed(self, contexts: List[str]) -> None:
        """为所有文本生成嵌入向量"""
        self.contexts = contexts
        for batch_context in batch_iterate(contexts, self.batch_size):
            batch_embeddings = self.generate_embedding(batch_context)
            self.embeddings.extend(batch_embeddings)
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.

from qdrant_client import QdrantClient, models

class QdrantVDB_QB:
    def __init__(self, collection_name: str, vector_dim: int = 768, batch_size: int = 512):
        self.collection_name = collection_name
        self.batch_size = batch_size
        self.vector_dim = vector_dim

    def define_client(self) -> None:
        """定义 Qdrant 客户端"""
        self.client = QdrantClient(url="http://localhost:6333", prefer_grpc=True)

    def create_collection(self) -> None:
        """创建 Qdrant 集合"""
        ifnot self.client.collection_exists(collection_name=self.collection_name):
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_cnotallow=models.VectorParams(
                    size=self.vector_dim,
                    distance=models.Distance.DOT,
                    on_disk=True
                ),
                optimizers_cnotallow=models.OptimizersConfigDiff(
                    default_segment_number=5,
                    indexing_threshold=0
                ),
                quantization_cnotallow=models.BinaryQuantization(
                    binary=models.BinaryQuantizationConfig(always_ram=True)
                )
            )

    def ingest_data(self, embeddata: EmbedData) -> None:
        """将嵌入数据导入 Qdrant"""
        for batch_context, batch_embeddings in zip(
            batch_iterate(embeddata.contexts, self.batch_size),
            batch_iterate(embeddata.embeddings, self.batch_size)
        ):
            self.client.upload_collection(
                collection_name=self.collection_name,
                vectors=batch_embeddings,
                payload=[{"context": context} for context in batch_context]
            )
        self.client.update_collection(
            collection_name=self.collection_name,
            optimizer_cnotallow=models.OptimizersConfigDiff(indexing_threshold=20000)
        )
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.

class Retriever:
    def __init__(self, vector_db: QdrantVDB_QB, embeddata: EmbedData):
        self.vector_db = vector_db
        self.embeddata = embeddata

    def search(self, query: str) -> List[dict]:
        """根据用户查询检索相关上下文"""
        query_embedding = self.embeddata.embed_model.get_query_embedding(query)
        result = self.vector_db.client.search(
            collection_name=self.vector_db.collection_name,
            query_vector=query_embedding,
            search_params=models.SearchParams(
                quantizatinotallow=models.QuantizationSearchParams(
                    ignore=False,
                    rescore=True,
                    oversampling=2.0
                )
            ),
            timeout=1000
        )
        return result
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.

from llama_index.llms.sambanovasystems import SambaNovaCloud
from llama_index.core.base.llms.types import ChatMessage, MessageRole

class RAG:
    def __init__(self, retriever: Retriever, llm_name: str = "Meta-Llama-3.1-405B-Instruct"):
        system_msg = ChatMessage(role=MessageRole.SYSTEM, cnotallow="You are a helpful assistant that answers questions about the user's document.")
        self.messages = [system_msg]
        self.llm_name = llm_name
        self.llm = self._setup_llm()
        self.retriever = retriever
        self.qa_prompt_tmpl_str = (
            "Context information is below.\n"
            "---------------------\n"
            "{context}\n"
            "---------------------\n"
            "Given the context information above I want you to think step by step to answer the query in a crisp manner, incase case you don't know the answer say 'I don't know!'.\n"
            "Query: {query}\n"
            "Answer: "
        )

    def _setup_llm(self) -> SambaNovaCloud:
        """设置 LLM"""
        return SambaNovaCloud(
            model=self.llm_name,
            temperature=0.7,
            context_window=100000
        )

    def generate_context(self, query: str) -> str:
        """生成上下文"""
        result = self.retriever.search(query)
        context = [dict(data) for data in result]
        combined_prompt = []
        for entry in context[:2]:
            combined_prompt.append(entry["payload"]["context"])
        return"\n\n---\n\n".join(combined_prompt)

    def query(self, query: str) -> str:
        """处理用户查询"""
        context = self.generate_context(query)
        prompt = self.qa_prompt_tmpl_str.format(cnotallow=context, query=query)
        user_msg = ChatMessage(role=MessageRole.USER, cnotallow=prompt)
        streaming_response = self.llm.stream_complete(user_msg.content)
        return streaming_response
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.

import assemblyai as aai
from typing import List, Dict

class Transcribe:
    def __init__(self, api_key: str):
        """初始化转录器"""
        aai.settings.api_key = api_key
        self.transcriber = aai.Transcriber()

    def transcribe_audio(self, audio_path: str) -> List[Dict[str, str]]:
        """转录音频文件"""
        config = aai.TranscriptionConfig(speaker_labels=True, speakers_expected=2)
        transcript = self.transcriber.transcribe(audio_path, cnotallow=config)
        speaker_transcripts = []
        for utterance in transcript.utterances:
            speaker_transcripts.append({
                "speaker": f"Speaker {utterance.speaker}",
                "text": utterance.text
            })
        return speaker_transcripts
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.

import os
import gc
import uuid
import tempfile
import base64
from dotenv import load_dotenv
import streamlit as st

# 加载环境变量
load_dotenv()

# 初始化会话状态
if"id"notin st.session_state:
    st.session_state.id = uuid.uuid4()
    st.session_state.file_cache = {}

# 侧边栏：上传音频文件
with st.sidebar:
    st.header("Add your audio file!")
    uploaded_file = st.file_uploader("Choose your audio file", type=["mp3", "wav", "m4a"])

    if uploaded_file:
        try:
            with tempfile.TemporaryDirectory() as temp_dir:
                file_path = os.path.join(temp_dir, uploaded_file.name)
                with open(file_path, "wb") as f:
                    f.write(uploaded_file.getvalue())
                file_key = f"{st.session_state.id}-{uploaded_file.name}"
                st.write("Transcribing with AssemblyAI and storing in vector database...")

                if file_key notin st.session_state.get('file_cache', {}):
                    # 初始化转录器
                    transcriber = Transcribe(api_key=os.getenv("ASSEMBLYAI_API_KEY"))
                    # 获取带有说话者标签的转录内容
                    transcripts = transcriber.transcribe_audio(file_path)
                    st.session_state.transcripts = transcripts

                    # 将每个说话者片段作为单独的文档进行嵌入
                    documents = [f"Speaker {t['speaker']}: {t['text']}"for t in transcripts]
                    embeddata = EmbedData(embed_model_name="BAAI/bge-large-en-v1.5", batch_size=32)
                    embeddata.embed(documents)

                    # 设置向量数据库
                    qdrant_vdb = QdrantVDB_QB(collection_name="chat_with_audios", batch_size=32, vector_dim=1024)
                    qdrant_vdb.define_client()
                    qdrant_vdb.create_collection()
                    qdrant_vdb.ingest_data(embeddata=embeddata)

                    # 设置检索器
                    retriever = Retriever(vector_db=qdrant_vdb, embeddata=embeddata)

                    # 设置 RAG 引擎
                    query_engine = RAG(retriever=retriever, llm_name="DeepSeek-R1-Distill-Llama-70B")
                    st.session_state.file_cache[file_key] = query_engine
                else:
                    query_engine = st.session_state.file_cache[file_key]

                # 提示用户文件已准备好
                st.success("Ready to Chat!")
                st.audio(uploaded_file)

                # 显示带有说话者标签的转录内容
                st.subheader("Transcript")
                with st.expander("Show full transcript", expanded=True):
                    for t in st.session_state.transcripts:
                        st.text(f"**{t['speaker']}**: {t['text']}")
        except Exception as e:
            st.error(f"An error occurred: {e}")
            st.stop()

# 初始化聊天历史
if"messages"notin st.session_state:
    st.session_state.messages = []

# 显示聊天历史
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 接受用户输入
if prompt := st.chat_input("Ask about the audio conversation..."):
    # 将用户消息添加到聊天历史中
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # 显示助手回应
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        streaming_response = query_engine.query(prompt)
        for chunk in streaming_response:
            try:
                new_text = chunk.raw["choices"][0]["delta"]["content"]
                full_response += new_text
                message_placeholder.markdown(full_response + "▌")
            except:
                pass
        message_placeholder.markdown(full_response)
    # 将助手回应添加到聊天历史中
    st.session_state.messages.append({"role": "assistant", "content": full_response})
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.