大模型应用系列：Query 变换的示例浅析

from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate

import os
from dotenv import load_dotenv

# Load environment variables from a .env file
load_dotenv()

# Set the OpenAI API key environment variable
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')

from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# LLM configuration
def get_rewriting_llm(model_name: str = "gpt-4o", temperature: float = 0, max_tokens: int = 4000) -> ChatOpenAI:
    return ChatOpenAI(temperature=temperature, model_name=model_name, max_tokens=max_tokens)

# Prompt Template for Query Rewriting
def create_query_rewrite_prompt() -> PromptTemplate:
    template = """
    You are an AI assistant tasked with reformulating user queries to improve retrieval in a RAG system. 
    Given the original query, rewrite it to be more specific, detailed, and likely to retrieve relevant information.

    Original query: {original_query}

    Rewritten query:
    """
    return PromptTemplate(input_variables=["original_query"], template=template)

# Build Query Rewriting Chain
def build_query_rewriting_chain(llm: ChatOpenAI) -> LLMChain:
    prompt_template = create_query_rewrite_prompt()
    return prompt_template | llm

# Function to Rewrite the Query
def rewrite_query(original_query: str, query_rewriter_chain: LLMChain) -> str:
    response = query_rewriter_chain.invoke(original_query)
    return response.content.strip()


if __name__ == "__main__":  
    llm = get_rewriting_llm()
    query_rewriter_chain = build_query_rewriting_chain(llm)
    original_query = "What are the impacts of climate change on the environment?"
    rewritten_query = rewrite_query(original_query, query_rewriter_chain)
    print("Original query:", original_query)
    print("\nRewritten query:", rewritten_query)

from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# LLM Configuration
def get_step_back_llm(model_name: str = "gpt-4o", temperature: float = 0, max_tokens: int = 4000) -> ChatOpenAI:
    return ChatOpenAI(temperature=temperature, model_name=model_name, max_tokens=max_tokens)

# Step-Back Prompt Template
def create_step_back_prompt_template() -> PromptTemplate:
    template = """
    You are an AI assistant tasked with generating broader, more general queries to improve context retrieval in a RAG system.
    Given the original query, generate a step-back query that is more general and can help retrieve relevant background information.

    Original query: {original_query}

    Step-back query:
    """
    return PromptTemplate(input_variables=["original_query"], template=template)

# Build Step-Back Query Chain
def build_step_back_chain(llm: ChatOpenAI) -> LLMChain:
    prompt_template = create_step_back_prompt_template()
    return prompt_template | llm

# Function to Generate Step-Back Query
def generate_step_back_query(original_query: str, step_back_chain: LLMChain) -> str:
    response = step_back_chain.invoke(original_query)
    return response.content.strip()

if __name__ == "__main__":
    llm = get_step_back_llm()
    step_back_chain = build_step_back_chain(llm)
    original_query = "What are the impacts of climate change on the environment?"
    step_back_query = generate_step_back_query(original_query, step_back_chain)
    print("Original query:", original_query)
    print("\nStep-back query:", step_back_query)

from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from typing import List

# LLM Configuration
def get_subquery_llm(model_name: str = "gpt-4o", temperature: float = 0, max_tokens: int = 4000) -> ChatOpenAI:
    return ChatOpenAI(temperature=temperature, model_name=model_name, max_tokens=max_tokens)

# Sub-query Decomposition Prompt Template
def create_subquery_decomposition_template() -> PromptTemplate:
    template = """
    You are an AI assistant tasked with breaking down complex queries into simpler sub-queries for a RAG system.
    Given the original query, decompose it into 2-4 simpler sub-queries that, when answered together, would provide a comprehensive response to the original query.

    Original query: {original_query}

    example: What are the impacts of climate change on the environment?

    Sub-queries:
    1. What are the impacts of climate change on biodiversity?
    2. How does climate change affect the oceans?
    3. What are the effects of climate change on agriculture?
    4. What are the impacts of climate change on human health?
    """
    return PromptTemplate(input_variables=["original_query"], template=template)

# Build Sub-query Decomposition Chain
def build_subquery_decomposer_chain(llm: ChatOpenAI) -> LLMChain:
    prompt_template = create_subquery_decomposition_template()
    return prompt_template | llm

# Function to Decompose Query into Sub-queries
def decompose_query(original_query: str, subquery_chain: LLMChain) -> List[str]:
    response = subquery_chain.invoke(original_query).content
    # Parse the sub-queries by splitting lines and removing unwanted text
    sub_queries = [q.strip() for q in response.split('\n') if q.strip() and not q.strip().startswith('Sub-queries:')]
    return sub_queries


if __name__ == "__main__":
    llm = get_subquery_llm()
    subquery_chain = build_subquery_decomposer_chain(llm)
    original_query = "What are the impacts of climate change on the environment?"
    sub_queries = decompose_query(original_query, subquery_chain)
    print("\nSub-queries:")
    for i, sub_query in enumerate(sub_queries, 1):
        print(f"{i}. {sub_query}")