阻止AI谎言：利用可信来源获得更智能的答案

import os
from pathlib import Path
import requests
import time
import json
from dotenv import load_dotenv
from llama_index.llms.openai import OpenAI
from llama_index.core.query_engine import CitationQueryEngine
from llama_index.core import (
    VectorStoreIndex,
    StorageContext,
    Settings,
    Document,
)
from llama_index.vector_stores.milvus import MilvusVectorStore

# Load environment variables
load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY")

# List of Wikipedia articles to scrape
WIKI_TITLES = [
    "Toronto",
    "Seattle",
    "San Francisco",
    "Chicago",
    "Boston",
    "Washington, D.C.",
    "Cambridge, Massachusetts",
    "Houston",
]

DATA_PATH = Path("data")
DATA_PATH.mkdir(exist_ok=True)

def scrape_wikipedia():
    # Scrape Wikipedia articles and save content with metadata.
    documents = []

    for title in WIKI_TITLES:
        response = requests.get(
            "https://en.wikipedia.org/w/api.php",
            params={
                "action": "query",
                "format": "json",
                "titles": title,
                "prop": "extracts",
                "explaintext": True,
            },
        ).json()
        page = next(iter(response["query"]["pages"].values()))
        wiki_text = page.get("extract", "")

        # Create metadata
        metadata = {
            "title": title,
            "source": "Wikipedia",
            "url": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}",
            "date_accessed": time.strftime("%Y-%m-%d"),
            "page_id": page.get("pageid", "Unknown"),
        }

        # Create Document object
        doc = Document(text=wiki_text, metadata=metadata)
        documents.append(doc)

        # Save content and metadata
        with open(DATA_PATH / f"{title}.txt", "w", encoding="utf-8") as f:
            f.write(wiki_text)

        with open(DATA_PATH / f"{title}_metadata.json", "w", encoding="utf-8") as f:
            json.dump(metadata, f, indent=2)

    return documents

def load_documents_with_metadata():
    # Load documents with their metadata from files.
    documents = []

    for content_file in DATA_PATH.glob("*.txt"):
        metadata_file = DATA_PATH / f"{content_file.stem}_metadata.json"

        # Load content and metadata
        with open(content_file, "r", encoding="utf-8") as f:
            content = f.read()

        metadata = (
            json.load(open(metadata_file, "r", encoding="utf-8"))
            if metadata_file.exists()
            else {
                "title": content_file.stem,
                "source": "Unknown",
                "url": "N/A",
                "date_accessed": "N/A",
            }
        )

        documents.append(Document(text=content, metadata=metadata))

    return documents

def setup_rag_with_citations():
    # Set up and run the RAG system with citations.
    try:
        # Configure settings
        Settings.llm = OpenAI(model="gpt-3.5-turbo", temperature=0)
        Settings.chunk_size = 512
        Settings.chunk_overlap = 50

        # Set up vector store with Milvus Lite
        vector_store = MilvusVectorStore(
            uri="./milvus_lite.db",  # Using local file for Milvus Lite
            collection_name="citations",
            dim=1536,
            overwrite=True,
        )

        # Load documents
        documents = load_documents_with_metadata()
        storage_context = StorageContext.from_defaults(vector_store=vector_store)

        # Create index
        index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

        # Create query engine
        query_engine = CitationQueryEngine.from_args(
            index,
            similarity_top_k=3,
            citation_chunk_size=512,
        )

        # Run query
        query = "Does Seattle or Houston have a bigger airport?"
        print(f"\nQuery: {query}\n")
        response = query_engine.query(query)

        # Print results
        print("Answer:")
        print("-" * 80)
        print(response.response)
        print("-" * 80)

        print_sources(response)

    except Exception as e:
        print(f"Error in RAG system: {e}")
        raise

def print_sources(response):
    # Print sources with enhanced metadata formatting.
    print("\nSources:\n")
    for idx, source in enumerate(response.source_nodes, start=1):
        metadata = source.node.metadata
        print(f"Source {idx}:")
        print(f"Title: {metadata.get('title', 'Unknown')}")
        print(f"Source: {metadata.get('source', 'Unknown')}")
        print(f"URL: {metadata.get('url', 'N/A')}")
        print(f"Accessed: {metadata.get('date_accessed', 'N/A')}")
        print("\nExcerpt:")
        print("-" * 80)
        print(source.node.get_text().strip())
        print("-" * 80)
        print()

def main():
    # Main execution function.
    try:
        print("Scraping Wikipedia articles...")
        scrape_wikipedia()

        print("Setting up RAG system...")
        setup_rag_with_citations()

    except Exception as e:
        print(f"An error occurred: {e}")


if __name__ == "__main__":
    main()