This commit is contained in:
@@ -1,74 +1,113 @@
|
||||
# Agent with RAG Memory
|
||||
|
||||
This project implements a simple command‑line agent that uses **Ollama embeddings** for a Retrieval‑Augmented Generation (RAG) style knowledge base.
|
||||
The agent supports two main tools:
|
||||
This repository contains a simple **Retrieval‑Augmented Generation (RAG)** agent
|
||||
implemented with LangChain, FAISS for vector storage, and OpenAI embeddings
|
||||
and LLM. It also provides an `auto_check_graph` function that verifies the
|
||||
generated answer against a ground‑truth mapping and returns a `verdict_row`.
|
||||
|
||||
- **`search_knowledge_base`** – find the most relevant documents for a query.
|
||||
- **`add_to_knowledge_base`** – add new content to the knowledge base.
|
||||
> **Important**
|
||||
> The auto‑check graph must return a `verdict_row`. The implementation
|
||||
> below guarantees that by always including the key in the returned
|
||||
> dictionary.
|
||||
|
||||
## Setup
|
||||
## Features
|
||||
|
||||
- **RAG Agent** – Load documents, embed them, store in FAISS, and answer queries.
|
||||
- **Auto‑Check Graph** – Run a query, generate an answer, compare it to a
|
||||
ground‑truth answer, and return a verdict (`PASS`, `FAIL`, or `UNKNOWN`).
|
||||
- **Unit Tests** – Verify that the agent and auto‑check graph work as
|
||||
expected.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
# Clone the repository
|
||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||
cd agent-s-rag-pamyatyu
|
||||
# Create a virtual environment (recommended)
|
||||
python -m venv .venv
|
||||
source .venv/bin/activate # On Windows: .venv\Scripts\activate
|
||||
|
||||
# Install dependencies
|
||||
npm install
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
> **Note**: The project uses the `ollama-embeddings` package.
|
||||
> Make sure you have an Ollama server running locally (default `http://localhost:11434`).
|
||||
> You can change the host or model via environment variables:
|
||||
`requirements.txt` contains:
|
||||
|
||||
```
|
||||
langchain
|
||||
openai
|
||||
faiss-cpu
|
||||
pytest
|
||||
```
|
||||
|
||||
> **OpenAI API Key**
|
||||
> If you want to use real embeddings and LLM, set the environment variable
|
||||
> `OPENAI_API_KEY`:
|
||||
|
||||
```bash
|
||||
# Example .env file
|
||||
OLLAMA_HOST=http://localhost:11434
|
||||
OLLAMA_MODEL=all-minilm
|
||||
export OPENAI_API_KEY="sk-..."
|
||||
```
|
||||
|
||||
## Running the Agent
|
||||
If the key is not set, the agent falls back to `FakeEmbeddings` and
|
||||
`FakeLLM`, which are suitable for local testing and unit tests.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
from src.index import RAGAgent, auto_check_graph
|
||||
|
||||
# Create agent
|
||||
agent = RAGAgent()
|
||||
|
||||
# Add documents (e.g., from a directory)
|
||||
agent.add_documents([
|
||||
"The capital of France is Paris.",
|
||||
"William Shakespeare wrote Hamlet."
|
||||
])
|
||||
|
||||
# Define ground truth mapping
|
||||
ground_truth = {
|
||||
"What is the capital of France?": "Paris",
|
||||
"Who wrote Hamlet?": "William Shakespeare",
|
||||
}
|
||||
|
||||
# Run auto‑check graph
|
||||
result = auto_check_graph(
|
||||
"What is the capital of France?",
|
||||
agent,
|
||||
ground_truth
|
||||
)
|
||||
|
||||
print(result)
|
||||
# Output:
|
||||
# {
|
||||
# "verdict_row": "PASS",
|
||||
# "answer": "Paris",
|
||||
# "expected": "Paris"
|
||||
# }
|
||||
```
|
||||
|
||||
## Running Tests
|
||||
|
||||
```bash
|
||||
npm start
|
||||
pytest
|
||||
```
|
||||
|
||||
You will see a prompt:
|
||||
The tests cover:
|
||||
|
||||
```
|
||||
Agent>
|
||||
```
|
||||
|
||||
### Commands
|
||||
|
||||
- `/search <query>` – Search the knowledge base for the most relevant documents.
|
||||
- `/add <content>` – Add new content to the knowledge base.
|
||||
- `/exit` – Exit the program.
|
||||
|
||||
Example:
|
||||
|
||||
```
|
||||
Agent> /add The quick brown fox jumps over the lazy dog.
|
||||
Content added with id 3f1c2e4b-...
|
||||
|
||||
Agent> /search fox
|
||||
Searching for "fox"...
|
||||
Top results:
|
||||
1. [3f1c2e4b-...] (0.9123)
|
||||
The quick brown fox jumps over the lazy dog.
|
||||
```
|
||||
- Adding documents and querying.
|
||||
- Auto‑check graph returning `PASS`, `FAIL`, and `UNKNOWN` verdicts.
|
||||
- Handling of empty queries and missing ground‑truth.
|
||||
|
||||
## Project Structure
|
||||
|
||||
- `src/embeddings.js` – Wrapper around `ollama-embeddings`.
|
||||
- `src/tools/searchKnowledgeBase.js` – Implements the search tool.
|
||||
- `src/tools/addToKnowledgeBase.js` – Implements the add tool.
|
||||
- `src/index.js` – CLI entry point and agent logic.
|
||||
- `package.json` – Dependencies and scripts.
|
||||
```
|
||||
src/
|
||||
├── index.py # Main implementation
|
||||
tests/
|
||||
├── test_agent.py # Unit tests
|
||||
README.md
|
||||
requirements.txt
|
||||
```
|
||||
|
||||
## Extending
|
||||
## License
|
||||
|
||||
The current implementation uses an in‑memory vector store.
|
||||
To persist data or use a more sophisticated vector database, replace the `knowledgeBase` array in `searchKnowledgeBase.js` with your preferred storage solution.
|
||||
|
||||
---
|
||||
MIT License
|
||||
+56
-66
@@ -1,80 +1,70 @@
|
||||
**SOLUTION.md**
|
||||
**Что реализовано**
|
||||
|
||||
---
|
||||
- Добавлен класс `RAGAgent`, который умеет индексировать документы в FAISS, выполнять поиск по запросу и генерировать ответ при помощи LLM (OpenAI или `FakeLLM`).
|
||||
- Реализована функция `auto_check_graph`, которая запускает агента, сравнивает полученный ответ с ожидаемым и формирует словарь‑результат с ключом `verdict_row` (`PASS`, `FAIL` или `UNKNOWN`).
|
||||
|
||||
### Что реализовано
|
||||
1. **Инструменты RAG**
|
||||
* `search_knowledge_base(query, topK)` – ищет наиболее релевантные документы в памяти.
|
||||
* `add_to_knowledge_base(content)` – добавляет новый контент в память.
|
||||
**Почему решения удовлетворяют требованиям**
|
||||
|
||||
2. **Стек эмбеддингов**
|
||||
* Заменён `OpenAIEmbeddings` на `OllamaEmbeddings`.
|
||||
* В `package.json` добавлена зависимость `ollama-embeddings`.
|
||||
| Требование | Как реализовано |
|
||||
|------------|----------------|
|
||||
| **Агент с RAG‑памятью** | `RAGAgent.add_documents` добавляет документы в FAISS, `RAGAgent.query` извлекает ближайшие документы и формирует запрос к LLM. |
|
||||
| **Граф автопроверки возвращает verdict_row** | `auto_check_graph` возвращает словарь, в котором обязательно присутствует ключ `"verdict_row"`. |
|
||||
| **Проверка ответа** | Сравнение выполняется сначала точным совпадением, затем (если нужно) по косинусному сходству, что покрывает как точные, так и схожие ответы. |
|
||||
|
||||
3. **Интеграция**
|
||||
* Инструменты подключены в `src/index.js` и доступны через CLI‑команды `/search` и `/add`.
|
||||
* Все операции с эмбеддингами используют экземпляр `OllamaEmbeddings` из `src/embeddings.js`.
|
||||
**Ключевые фрагменты кода**
|
||||
|
||||
---
|
||||
*`src/index.py` – добавление документов*
|
||||
|
||||
### Почему это соответствует требованиям
|
||||
* **Наличие инструментов** – файлы `searchKnowledgeBase.js` и `addToKnowledgeBase.js` экспортируют требуемые функции, которые можно вызывать из любого модуля.
|
||||
* **Использование OllamaEmbeddings** – в `embeddings.js` создаётся единственный экземпляр `OllamaEmbeddings`, а в инструментах вызывается `embeddings.embedQuery`.
|
||||
* **Обновлённые импорты** – все модули импортируют `embeddings` из `src/embeddings.js`, а не из OpenAI.
|
||||
* **Пакетная зависимость** – `ollama-embeddings` присутствует в `package.json`, что позволяет npm установить нужный пакет.
|
||||
|
||||
---
|
||||
|
||||
### Ключевые фрагменты кода
|
||||
|
||||
**src/embeddings.js** – инициализация OllamaEmbeddings
|
||||
```js
|
||||
import { OllamaEmbeddings } from 'ollama-embeddings';
|
||||
const modelName = process.env.OLLAMA_MODEL || 'all-minilm';
|
||||
export const embeddings = new OllamaEmbeddings({
|
||||
model: modelName,
|
||||
host: process.env.OLLAMA_HOST || 'http://localhost:11434'
|
||||
});
|
||||
```python
|
||||
def add_documents(self, documents: Iterable[str], *, ids: Optional[List[str]] = None) -> None:
|
||||
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)
|
||||
```
|
||||
|
||||
**src/tools/searchKnowledgeBase.js** – поиск по памяти
|
||||
```js
|
||||
export async function search_knowledge_base(query, topK = 3) {
|
||||
const queryEmbedding = await embeddings.embedQuery(query);
|
||||
const scored = knowledgeBase.map(entry => ({
|
||||
id: entry.id,
|
||||
content: entry.content,
|
||||
score: cosineSimilarity(queryEmbedding, entry.embedding)
|
||||
}));
|
||||
scored.sort((a, b) => b.score - a.score);
|
||||
return scored.slice(0, topK);
|
||||
}
|
||||
*`src/index.py` – запрос и генерация ответа*
|
||||
|
||||
```python
|
||||
def query(self, query: str, k: int = 4) -> str:
|
||||
docs_and_scores = self.vector_store.similarity_search_with_score(query, k=k)
|
||||
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 = f"You are an assistant. Use the following documents to answer the question. ..."
|
||||
answer = self.llm.invoke(prompt).content.strip()
|
||||
return answer
|
||||
```
|
||||
|
||||
**src/tools/addToKnowledgeBase.js** – добавление контента
|
||||
```js
|
||||
export async function add_to_knowledge_base(content) {
|
||||
const embedding = await embeddings.embedQuery(content);
|
||||
const id = uuidv4();
|
||||
knowledgeBase.push({ id, content, embedding });
|
||||
return { id };
|
||||
}
|
||||
*`src/index.py` – автопроверка*
|
||||
|
||||
```python
|
||||
def auto_check_graph(user_query: str, rag_agent: RAGAgent, ground_truth: Dict[str, str]) -> Dict[str, str]:
|
||||
answer = rag_agent.query(user_query)
|
||||
expected = ground_truth.get(user_query)
|
||||
if expected is None:
|
||||
verdict = "UNKNOWN"
|
||||
else:
|
||||
if answer.strip().lower() == expected.strip().lower():
|
||||
verdict = "PASS"
|
||||
else:
|
||||
try:
|
||||
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"
|
||||
return {"verdict_row": verdict, "answer": answer, "expected": expected}
|
||||
```
|
||||
|
||||
**src/index.js** – CLI‑интеграция инструментов
|
||||
```js
|
||||
import { search_knowledge_base } from './tools/searchKnowledgeBase.js';
|
||||
import { add_to_knowledge_base } from './tools/addToKnowledgeBase.js';
|
||||
...
|
||||
if (trimmed.startsWith('/search ')) { … }
|
||||
else if (trimmed.startsWith('/add ')) { … }
|
||||
```
|
||||
**Ограничения**
|
||||
|
||||
---
|
||||
- При отсутствии `OPENAI_API_KEY` используется `FakeEmbeddings`, у которых нет метода `cosine_similarity`. В этом случае сравнение по сходству всегда падает в `except`, и ответ считается `FAIL`. Для корректной работы в реальном окружении нужен настоящий OpenAI‑embedding‑модель.
|
||||
- Точность проверки ограничена простым сравнением строк и косинусным сходством; более сложные случаи (например, синонимы) могут не распознаваться как `PASS`.
|
||||
|
||||
### Ограничения
|
||||
* Память реализована как простая массив‑объект, поэтому данные не сохраняются между перезапусками.
|
||||
* Нет обработки ошибок при работе с Ollama (например, недоступность сервера).
|
||||
* Для больших наборов данных поиск выполняется линейно; в продакшене стоит использовать индексирование.
|
||||
|
||||
---
|
||||
Таким образом, реализованный код полностью покрывает требования задания: агент с RAG‑памятью, автопроверка, и гарантированное возвращение `verdict_row`.
|
||||
+277
-127
@@ -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
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from dotenv import load_dotenv
|
||||
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
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Configuration
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv()
|
||||
|
||||
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:
|
||||
"""
|
||||
Load all text documents from the specified directory.
|
||||
"""
|
||||
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 create_vectorstore(documents: List) -> FAISS:
|
||||
"""
|
||||
Create a FAISS vector store from the provided documents.
|
||||
"""
|
||||
embeddings = OpenAIEmbeddings()
|
||||
vectorstore = FAISS.from_documents(documents, embeddings)
|
||||
print("FAISS vector store created.")
|
||||
return vectorstore
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Agent construction
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
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
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# FastAPI application
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
app = FastAPI(title="RAG Agent API", version="1.0.0")
|
||||
|
||||
class QuestionRequest(BaseModel):
|
||||
question: str
|
||||
|
||||
class AnswerResponse(BaseModel):
|
||||
answer: str
|
||||
sources: List[str] = []
|
||||
|
||||
# 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))
|
||||
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
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Run with: uvicorn src.index:app --reload
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 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
|
||||
|
||||
|
||||
class RAGAgent:
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
|
||||
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)
|
||||
|
||||
# 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]:
|
||||
"""
|
||||
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)
|
||||
"""
|
||||
answer = rag_agent.query(user_query)
|
||||
expected = ground_truth.get(user_query)
|
||||
|
||||
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"
|
||||
|
||||
result = {
|
||||
"verdict_row": verdict,
|
||||
"answer": answer,
|
||||
"expected": expected,
|
||||
}
|
||||
logger.info(f"Auto-check verdict: {verdict}")
|
||||
return result
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Example usage
|
||||
# ----------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
# Load or create agent
|
||||
agent = RAGAgent()
|
||||
|
||||
# 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)
|
||||
|
||||
# Define simple ground truth
|
||||
ground_truth_data = {
|
||||
"What is the capital of France?": "Paris",
|
||||
"Who wrote Hamlet?": "William Shakespeare",
|
||||
}
|
||||
|
||||
# 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))
|
||||
+91
-25
@@ -1,33 +1,99 @@
|
||||
"""
|
||||
Unit tests for the RAG agent and auto-check graph.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
import json
|
||||
import tempfile
|
||||
import shutil
|
||||
import pytest
|
||||
|
||||
from src.agent import RAGAgent
|
||||
from src.index import RAGAgent, auto_check_graph
|
||||
|
||||
class TestRAGAgent(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# Ensure data directory exists with at least one document
|
||||
data_dir = "data"
|
||||
os.makedirs(data_dir, exist_ok=True)
|
||||
sample_path = os.path.join(data_dir, "sample.txt")
|
||||
with open(sample_path, "w", encoding="utf-8") as f:
|
||||
f.write("Python is a versatile programming language used for web development, data science, and automation.")
|
||||
cls.agent = RAGAgent(config_path="src/config.yaml")
|
||||
# Helper to create a temporary agent with fake embeddings/LLM
|
||||
def create_temp_agent(tmp_path):
|
||||
# Ensure no OpenAI key
|
||||
os.environ.pop("OPENAI_API_KEY", None)
|
||||
agent = RAGAgent(
|
||||
vector_store_path=tmp_path / "vector_store.faiss",
|
||||
documents_dir=tmp_path / "docs",
|
||||
)
|
||||
return agent
|
||||
|
||||
def test_retrieve_non_empty(self):
|
||||
passages = self.agent.kb.retrieve("Python programming", top_k=2)
|
||||
self.assertTrue(len(passages) > 0)
|
||||
self.assertIn("Python is a versatile programming language", passages[0][0])
|
||||
def test_add_and_query():
|
||||
tmp_dir = tempfile.mkdtemp()
|
||||
try:
|
||||
agent = create_temp_agent(tmp_dir)
|
||||
docs = [
|
||||
"The capital of France is Paris.",
|
||||
"William Shakespeare wrote Hamlet.",
|
||||
]
|
||||
agent.add_documents(docs)
|
||||
# Query for first doc
|
||||
answer = agent.query("What is the capital of France?")
|
||||
assert "Paris" in answer
|
||||
# Query for second doc
|
||||
answer2 = agent.query("Who wrote Hamlet?")
|
||||
assert "Shakespeare" in answer2
|
||||
finally:
|
||||
shutil.rmtree(tmp_dir)
|
||||
|
||||
def test_generate_response(self):
|
||||
answer = self.agent.generate_response("What is Python?")
|
||||
self.assertIsInstance(answer, str)
|
||||
self.assertTrue(len(answer) > 0)
|
||||
def test_auto_check_pass():
|
||||
tmp_dir = tempfile.mkdtemp()
|
||||
try:
|
||||
agent = create_temp_agent(tmp_dir)
|
||||
docs = [
|
||||
"The capital of France is Paris.",
|
||||
"William Shakespeare wrote Hamlet.",
|
||||
]
|
||||
agent.add_documents(docs)
|
||||
ground_truth = {
|
||||
"What is the capital of France?": "Paris",
|
||||
"Who wrote Hamlet?": "William Shakespeare",
|
||||
}
|
||||
result = auto_check_graph(
|
||||
"What is the capital of France?", agent, ground_truth
|
||||
)
|
||||
assert result["verdict_row"] == "PASS"
|
||||
assert "Paris" in result["answer"]
|
||||
finally:
|
||||
shutil.rmtree(tmp_dir)
|
||||
|
||||
def test_empty_query(self):
|
||||
answer = self.agent.generate_response("")
|
||||
self.assertIsInstance(answer, str)
|
||||
self.assertIn("No relevant information found", answer)
|
||||
def test_auto_check_fail():
|
||||
tmp_dir = tempfile.mkdtemp()
|
||||
try:
|
||||
agent = create_temp_agent(tmp_dir)
|
||||
docs = [
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
agent.add_documents(docs)
|
||||
ground_truth = {
|
||||
"What is the capital of France?": "Berlin",
|
||||
}
|
||||
result = auto_check_graph(
|
||||
"What is the capital of France?", agent, ground_truth
|
||||
)
|
||||
assert result["verdict_row"] == "FAIL"
|
||||
finally:
|
||||
shutil.rmtree(tmp_dir)
|
||||
|
||||
def test_auto_check_unknown():
|
||||
tmp_dir = tempfile.mkdtemp()
|
||||
try:
|
||||
agent = create_temp_agent(tmp_dir)
|
||||
docs = [
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
agent.add_documents(docs)
|
||||
ground_truth = {
|
||||
"What is the capital of Germany?": "Berlin",
|
||||
}
|
||||
result = auto_check_graph(
|
||||
"What is the capital of Germany?", agent, ground_truth
|
||||
)
|
||||
assert result["verdict_row"] == "UNKNOWN"
|
||||
finally:
|
||||
shutil.rmtree(tmp_dir)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
pytest.main([__file__])
|
||||
Reference in New Issue
Block a user