This commit is contained in:
+263
-113
@@ -1,139 +1,289 @@
|
||||
"""
|
||||
Agent with Retrieval-Augmented Generation (RAG) memory.
|
||||
RAG Agent with Auto-Check Graph
|
||||
================================
|
||||
|
||||
This module implements a FastAPI application that exposes a single endpoint
|
||||
`/ask` for querying an RAG-enabled agent. The agent uses LangChain to
|
||||
embed documents from a local `data/` directory into a FAISS vector store,
|
||||
retrieves relevant passages for a user query, and generates a response
|
||||
using OpenAI's GPT-4 model.
|
||||
|
||||
Prerequisites:
|
||||
- Python 3.11+
|
||||
- OpenAI API key set in the environment variable `OPENAI_API_KEY`
|
||||
(or in a `.env` file in the project root).
|
||||
- Text files placed in the `data/` directory (one file per document).
|
||||
This module implements a simple Retrieval-Augmented Generation (RAG) agent
|
||||
using LangChain, FAISS for vector storage, and OpenAI embeddings and
|
||||
LLM. It also provides an `auto_check_graph` function that runs a
|
||||
verification routine against a ground‑truth answer and returns a
|
||||
`verdict_row` indicating whether the generated answer matches the
|
||||
expected answer.
|
||||
|
||||
Author: Artur Kuzakhmetov
|
||||
Version: 20
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from dotenv import load_dotenv
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterable, List, Optional
|
||||
|
||||
# LangChain imports
|
||||
from langchain.document_loaders import DirectoryLoader
|
||||
from langchain.embeddings.openai import OpenAIEmbeddings
|
||||
from langchain.vectorstores import FAISS
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.llms import OpenAI
|
||||
try:
|
||||
from langchain.embeddings.openai import OpenAIEmbeddings
|
||||
from langchain.embeddings.fake import FakeEmbeddings
|
||||
from langchain.llms.openai import ChatOpenAI
|
||||
from langchain.llms.fake import FakeLLM
|
||||
from langchain.vectorstores.faiss import FAISS
|
||||
from langchain.docstore.document import Document
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Required LangChain packages are missing. "
|
||||
"Install with: pip install langchain openai faiss-cpu"
|
||||
) from exc
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Configuration
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv()
|
||||
# Default constants
|
||||
DEFAULT_VECTOR_STORE_PATH = Path("vector_store.faiss")
|
||||
DEFAULT_DOCUMENTS_DIR = Path("documents")
|
||||
DEFAULT_EMBEDDING_MODEL = "text-embedding-ada-002"
|
||||
DEFAULT_LLM_MODEL = "gpt-3.5-turbo"
|
||||
SIMILARITY_THRESHOLD = 0.8 # Cosine similarity threshold for PASS
|
||||
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
sys.exit("Error: OPENAI_API_KEY not found in environment variables.")
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Data loading and vector store initialization
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
DATA_DIR = Path(__file__).parent.parent / "data"
|
||||
|
||||
def load_documents(path: Path) -> List:
|
||||
class RAGAgent:
|
||||
"""
|
||||
Load all text documents from the specified directory.
|
||||
Retrieval-Augmented Generation (RAG) agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
embedding_model : str, optional
|
||||
Name of the OpenAI embedding model to use. If the
|
||||
``OPENAI_API_KEY`` environment variable is not set, a
|
||||
``FakeEmbeddings`` instance is used.
|
||||
llm_model : str, optional
|
||||
Name of the OpenAI LLM to use. If the ``OPENAI_API_KEY`` is
|
||||
not set, a ``FakeLLM`` instance is used.
|
||||
vector_store_path : Path, optional
|
||||
Path to the FAISS vector store file.
|
||||
documents_dir : Path, optional
|
||||
Directory containing text files to be indexed.
|
||||
"""
|
||||
if not path.exists() or not path.is_dir():
|
||||
print(f"Warning: Data directory '{path}' not found. No documents loaded.")
|
||||
return []
|
||||
|
||||
loader = DirectoryLoader(str(path), glob="**/*.txt")
|
||||
documents = loader.load()
|
||||
print(f"Loaded {len(documents)} documents from '{path}'.")
|
||||
return documents
|
||||
def __init__(
|
||||
self,
|
||||
embedding_model: str = DEFAULT_EMBEDDING_MODEL,
|
||||
llm_model: str = DEFAULT_LLM_MODEL,
|
||||
vector_store_path: Path = DEFAULT_VECTOR_STORE_PATH,
|
||||
documents_dir: Path = DEFAULT_DOCUMENTS_DIR,
|
||||
) -> None:
|
||||
self.embedding_model_name = embedding_model
|
||||
self.llm_model_name = llm_model
|
||||
self.vector_store_path = Path(vector_store_path)
|
||||
self.documents_dir = Path(documents_dir)
|
||||
|
||||
def create_vectorstore(documents: List) -> FAISS:
|
||||
# Initialize embeddings
|
||||
if os.getenv("OPENAI_API_KEY"):
|
||||
self.embeddings = OpenAIEmbeddings(
|
||||
model=self.embedding_model_name,
|
||||
chunk_size=512,
|
||||
)
|
||||
self.llm = ChatOpenAI(
|
||||
model=self.llm_model_name,
|
||||
temperature=0.0,
|
||||
)
|
||||
logger.info("Using OpenAI embeddings and LLM.")
|
||||
else:
|
||||
# Fallback for local testing
|
||||
self.embeddings = FakeEmbeddings()
|
||||
self.llm = FakeLLM()
|
||||
logger.warning(
|
||||
"OPENAI_API_KEY not found. Using FakeEmbeddings and FakeLLM."
|
||||
)
|
||||
|
||||
# Load or create vector store
|
||||
if self.vector_store_path.exists():
|
||||
self.vector_store = FAISS.load_local(
|
||||
self.vector_store_path,
|
||||
self.embeddings,
|
||||
allow_dangerous_deserialization=True,
|
||||
)
|
||||
logger.info(
|
||||
f"Loaded existing vector store from {self.vector_store_path}"
|
||||
)
|
||||
else:
|
||||
self.vector_store = FAISS(
|
||||
embedding_function=self.embeddings,
|
||||
index=None,
|
||||
)
|
||||
logger.info("Created new empty vector store.")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Document management
|
||||
# ------------------------------------------------------------------
|
||||
def add_documents(
|
||||
self,
|
||||
documents: Iterable[str],
|
||||
*,
|
||||
ids: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add a collection of documents to the vector store.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
documents : Iterable[str]
|
||||
Text content of documents to add.
|
||||
ids : List[str], optional
|
||||
Optional list of identifiers for the documents.
|
||||
"""
|
||||
docs = [
|
||||
Document(page_content=doc, metadata={"id": doc_id})
|
||||
for doc, doc_id in zip(documents, ids or [None] * len(documents))
|
||||
]
|
||||
self.vector_store.add_documents(docs)
|
||||
self.vector_store.save_local(self.vector_store_path)
|
||||
logger.info(f"Added {len(docs)} documents to vector store.")
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""
|
||||
Remove the persisted vector store file.
|
||||
"""
|
||||
if self.vector_store_path.exists():
|
||||
self.vector_store_path.unlink()
|
||||
logger.info(f"Deleted vector store file {self.vector_store_path}.")
|
||||
else:
|
||||
logger.info("No vector store file to delete.")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Querying
|
||||
# ------------------------------------------------------------------
|
||||
def query(self, query: str, k: int = 4) -> str:
|
||||
"""
|
||||
Retrieve relevant documents and generate an answer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The user query.
|
||||
k : int, optional
|
||||
Number of nearest neighbors to retrieve.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Generated answer.
|
||||
"""
|
||||
if not query.strip():
|
||||
logger.warning("Empty query received.")
|
||||
return "No query provided."
|
||||
|
||||
# Retrieve relevant documents
|
||||
docs_and_scores = self.vector_store.similarity_search_with_score(
|
||||
query, k=k
|
||||
)
|
||||
if not docs_and_scores:
|
||||
logger.info("No relevant documents found.")
|
||||
return "I couldn't find any relevant information."
|
||||
|
||||
# Build context string
|
||||
context = "\n\n".join(
|
||||
f"Document {i+1} (score={score:.3f}):\n{doc.page_content}"
|
||||
for i, (doc, score) in enumerate(docs_and_scores)
|
||||
)
|
||||
|
||||
# Prompt for LLM
|
||||
prompt = (
|
||||
f"You are an assistant. Use the following documents to answer the "
|
||||
f"question. If you cannot answer, say so.\n\n"
|
||||
f"Documents:\n{context}\n\n"
|
||||
f"Question: {query}\nAnswer:"
|
||||
)
|
||||
|
||||
# Generate answer
|
||||
answer = self.llm.invoke(prompt).content.strip()
|
||||
logger.info(f"Generated answer for query: {query}")
|
||||
return answer
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Auto-check graph
|
||||
# ----------------------------------------------------------------------
|
||||
def auto_check_graph(
|
||||
user_query: str,
|
||||
rag_agent: RAGAgent,
|
||||
ground_truth: Dict[str, str],
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
Create a FAISS vector store from the provided documents.
|
||||
Run the RAG agent on a query and verify the answer against a
|
||||
ground‑truth mapping.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
user_query : str
|
||||
The query to process.
|
||||
rag_agent : RAGAgent
|
||||
Instance of the RAG agent.
|
||||
ground_truth : Dict[str, str]
|
||||
Mapping from query to expected answer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Dict[str, str]
|
||||
Dictionary containing:
|
||||
- verdict_row: 'PASS', 'FAIL', or 'UNKNOWN'
|
||||
- answer: Generated answer
|
||||
- expected: Expected answer (may be None)
|
||||
"""
|
||||
embeddings = OpenAIEmbeddings()
|
||||
vectorstore = FAISS.from_documents(documents, embeddings)
|
||||
print("FAISS vector store created.")
|
||||
return vectorstore
|
||||
answer = rag_agent.query(user_query)
|
||||
expected = ground_truth.get(user_query)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Agent construction
|
||||
# --------------------------------------------------------------------------- #
|
||||
if expected is None:
|
||||
verdict = "UNKNOWN"
|
||||
else:
|
||||
# Simple exact match check
|
||||
if answer.strip().lower() == expected.strip().lower():
|
||||
verdict = "PASS"
|
||||
else:
|
||||
# Fallback similarity check using embeddings
|
||||
try:
|
||||
# Use the same embeddings as the agent
|
||||
query_vec = rag_agent.embeddings.embed_query(user_query)
|
||||
answer_vec = rag_agent.embeddings.embed_query(answer)
|
||||
similarity = rag_agent.embeddings.cosine_similarity(
|
||||
query_vec, answer_vec
|
||||
)
|
||||
verdict = "PASS" if similarity >= SIMILARITY_THRESHOLD else "FAIL"
|
||||
except Exception as exc:
|
||||
logger.warning(f"Similarity check failed: {exc}")
|
||||
verdict = "FAIL"
|
||||
|
||||
def build_agent(vectorstore: FAISS) -> RetrievalQA:
|
||||
"""
|
||||
Build a RetrievalQA chain that uses the vector store for retrieval
|
||||
and OpenAI GPT-4 for generation.
|
||||
"""
|
||||
llm = OpenAI(model_name="gpt-4", temperature=0, openai_api_key=OPENAI_API_KEY)
|
||||
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
|
||||
qa_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=retriever,
|
||||
return_source_documents=True,
|
||||
)
|
||||
print("RetrievalQA agent constructed.")
|
||||
return qa_chain
|
||||
result = {
|
||||
"verdict_row": verdict,
|
||||
"answer": answer,
|
||||
"expected": expected,
|
||||
}
|
||||
logger.info(f"Auto-check verdict: {verdict}")
|
||||
return result
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# FastAPI application
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
app = FastAPI(title="RAG Agent API", version="1.0.0")
|
||||
# ----------------------------------------------------------------------
|
||||
# Example usage
|
||||
# ----------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
# Load or create agent
|
||||
agent = RAGAgent()
|
||||
|
||||
class QuestionRequest(BaseModel):
|
||||
question: str
|
||||
# Example: add documents from a directory
|
||||
if agent.documents_dir.exists():
|
||||
docs = []
|
||||
for file_path in agent.documents_dir.glob("*.txt"):
|
||||
docs.append(file_path.read_text(encoding="utf-8"))
|
||||
if docs:
|
||||
agent.add_documents(docs)
|
||||
|
||||
class AnswerResponse(BaseModel):
|
||||
answer: str
|
||||
sources: List[str] = []
|
||||
# Define simple ground truth
|
||||
ground_truth_data = {
|
||||
"What is the capital of France?": "Paris",
|
||||
"Who wrote Hamlet?": "William Shakespeare",
|
||||
}
|
||||
|
||||
# Global variables to hold the agent and vector store
|
||||
vectorstore: FAISS = None
|
||||
agent: RetrievalQA = None
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup_event():
|
||||
"""
|
||||
Load documents, create vector store, and build the agent on startup.
|
||||
"""
|
||||
global vectorstore, agent
|
||||
docs = load_documents(DATA_DIR)
|
||||
vectorstore = create_vectorstore(docs)
|
||||
agent = build_agent(vectorstore)
|
||||
|
||||
@app.post("/ask", response_model=AnswerResponse)
|
||||
def ask_question(request: QuestionRequest):
|
||||
"""
|
||||
Endpoint to query the RAG agent.
|
||||
"""
|
||||
if not agent:
|
||||
raise HTTPException(status_code=500, detail="Agent not initialized.")
|
||||
try:
|
||||
result = agent({"question": request.question})
|
||||
answer = result.get("answer", "")
|
||||
sources = [doc.metadata.get("source", "") for doc in result.get("source_documents", [])]
|
||||
return AnswerResponse(answer=answer, sources=sources)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Run with: uvicorn src.index:app --reload
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Run auto-check graph
|
||||
query = "What is the capital of France?"
|
||||
result = auto_check_graph(query, agent, ground_truth_data)
|
||||
print(json.dumps(result, indent=2))
|
||||
Reference in New Issue
Block a user