RAG 知识库：增强式信息检索生成技术解析

import re
import numpy as np
from openai import OpenAI
import json
import os

class APIClient:
    def __init__(self, config_file):
        """
        初始化 API 客户端，从配置文件中加载必要信息。
        """
        with open(config_file, 'r', encoding='utf-8') as f:
            self.config = json.load(f)
            self.client = OpenAI(
                api_key=self.config["api_key"],
                base_url=self.config["base_url"]
            )
            self.provider = "openai"
        

    def get_embeddings(self, texts):
        """
        获取文本的嵌入向量。
        """
        response = self.client.embeddings.create(
            input=texts,
            model=self.config["embedding_model"]
        )
        return [item.embedding for item in response.data]

    def generate_response(self, prompt, context):
        """
        生成回答。
        """
        response = self.client.chat.completions.create(
            model=self.config["llm_model"],
            messages=[
                {"role": "system", "content": "你是个人助理，请根据我给的内容，解答我的问题，不要原创。"},
                {"role": "user", "content": f"{prompt}\n\n 。相关知识如下，请根据下列相关知识回答，不要自由发挥: {context}"}
            ],
            max_tokens=2000
        )
        return response.choices[0].message.content

def save_embeddings_to_json(embeddings, file_path):
    with open(file_path, 'w', encoding='utf-8') as f:
        json.dump(embeddings, f, ensure_ascii=False, indent=4)

def load_embeddings_from_json(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        embeddings = json.load(f)
    return embeddings          

def read_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
    return content

def split_text_into_chunks(text, chunk_size=2000):
    sentences = re.split(r'(?<=[。！？])', text)
    chunks = []
    current_chunk = ""
    
    for sentence in sentences:
        if len(current_chunk) + len(sentence) <= chunk_size:
            current_chunk += sentence
        else:
            chunks.append(current_chunk)
            current_chunk = sentence
    
    if current_chunk:
        chunks.append(current_chunk)
    
    return chunks

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def main(api_provider):
    # 根据选择的 API 提供商加载配置文件
    if api_provider == "openai":
        config_file = "openai_config.json"
    elif api_provider == "ali":
        config_file = "ali_config.json"
    else:
        raise ValueError("Unsupported API provider.")

    # 初始化 API 客户端
    api_client = APIClient(config_file)    

    file_path = 'sub.md'
    text = read_file(file_path)
    chunks = split_text_into_chunks(text)
    print(f"len of chunks {len(chunks)}, last {chunks[-1]}")
    
    embeddings_file = 'embeddings.json'
    if os.path.exists(embeddings_file):
        print(f"文件 {embeddings_file} 存在")
        embeddings = load_embeddings_from_json(embeddings_file)
    else:
        print(f"文件 {embeddings_file} 不存在")
        embeddings = api_client.get_embeddings(chunks)
        save_embeddings_to_json(embeddings, 'embeddings.json')

    query = "我想喝南瓜汤，找篇菜谱给我"
    query_embedding = api_client.get_embeddings([query])[0]
    
    similarities = [cosine_similarity(query_embedding, chunk_embedding) for chunk_embedding in embeddings]
    most_similar_chunk_index = np.argmax(similarities)
    
    context = chunks[most_similar_chunk_index]
    print(f"\n\n\n\n相关内容: {context}\n\n\n\n")
    response = api_client.generate_response(query, context)
    
    print("\n\nGenerated Response:\n\n")
    print(response)

if __name__ == "__main__":
    # 选择 API 提供商
    api_provider = "ali"  # 可选 "openai" 或 "ali"
    main(api_provider)