动态RAG应用在Langserve中实现动态输入教程

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableParallel
from langchain_core.output_parsers import StrOutputParser
from langchain_community.vectorstores import FAISS # 示例，实际可替换为其他向量库
from langchain_community.embeddings import OpenAIEmbeddings # 示例，实际可替换为其他嵌入模型
from langchain_community.document_loaders import TextLoader # 示例
from langchain.text_splitter import CharacterTextSplitter # 示例
from langchain_openai import ChatOpenAI # 示例，实际可替换为其他LLM
from operator import itemgetter
import os

# 假设已经配置了OPENAI_API_KEY环境变量
# os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"

# 1. 模拟数据加载和索引
# 实际应用中，这里会加载你的文档并创建向量存储
# 为了示例，我们创建一个简单的内存向量存储
with open("example_docs.txt", "w") as f:
    f.write("Finance accounts typically contain records of financial transactions, including income, expenses, assets, liabilities, and equity. They are used to track a company's financial performance and position.\n")
    f.write("A balance sheet provides a snapshot of a company's assets, liabilities, and owner's equity at a specific point in time.\n")
    f.write("The income statement reports a company's financial performance over a specific accounting period, showing revenues and expenses.\n")

loader = TextLoader("example_docs.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()
vectorstore = FAISS.from_documents(docs, embeddings)
retriever_llm = vectorstore.as_retriever() # 注意：这里命名为retriever_llm是为了与原问题保持一致，实际是retriever

# 2. 初始化LLM
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

# 3. 定义提示模板
# 提示模板将接受 'context', 'question' 和 'lang' 作为输入
template = """根据以下上下文信息回答问题，并严格用{lang}语言输出：
{context}

问题: {question}
答案:"""
prompt = ChatPromptTemplate.from_template(template)

# 辅助函数：格式化检索到的文档
def format_docs(docs):
    return "\n\n".join([doc.page_content for doc in docs])

# 定义RAG链
rag_chain = (
    # 1. 并行处理输入：
    #    - 'context': 从输入中提取 'question'，传递给检索器，然后格式化文档。
    #    - 'question': 直接从输入中提取 'question'。
    #    - 'lang': 直接从输入中提取 'lang'。
    RunnableParallel({
        "context": itemgetter("question") | retriever_llm | format_docs,
        "question": itemgetter("question"),
        "lang": itemgetter("lang")
    })
    | prompt # 2. 将处理后的输入传递给提示模板
    | llm    # 3. 将提示模板的输出传递给LLM
    | StrOutputParser() # 4. 解析LLM的输出为字符串
)

from fastapi import FastAPI
from langserve import add_routes
from pydantic import BaseModel, Field

# 定义Langserve的输入模型
class InputQuestion(BaseModel):
    question: str = Field(..., description="The user's query for the RAG system.")
    lang: str = Field("English", description="The desired output language (e.g., 'English', 'Chinese').")

app = FastAPI(
    title="Dynamic RAG Langserve Application",
    version="1.0",
    description="A RAG application with dynamic question and language inputs."
)

# 添加路由
# input_type 参数确保Langserve知道如何解析传入的JSON请求体
add_routes(
    app,
    rag_chain,
    path="/dynamic-rag",
    input_type=InputQuestion, # 指定输入模型
    # output_type=str # 如果需要，可以指定输出类型，默认通常是字符串
)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="localhost", port=8000)