This commit is contained in:
+2
-3
@@ -1,3 +1,2 @@
|
||||
# OpenAI API key (optional). If not set, the LLM will use a simple echo fallback.
|
||||
OPENAI_API_KEY=
|
||||
PORT=3000
|
||||
# Copy this file to .env and replace with your actual OpenAI API key
|
||||
OPENAI_API_KEY=sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
|
||||
@@ -1,100 +1,166 @@
|
||||
# Educational Agent with Retrieval-Augmented Generation (RAG) Memory
|
||||
# RAG Agent with Retrieval-Augmented Generation
|
||||
|
||||
This repository contains a lightweight educational agent that demonstrates
|
||||
Retrieval-Augmented Generation (RAG) using only Python standard libraries.
|
||||
The implementation is fully self‑contained and does not rely on external
|
||||
AI services or heavy dependencies, making it suitable for the Deep Agents
|
||||
Virtual File System environment.
|
||||
**Version:** 20
|
||||
**Author:** Artur Kuzakhmetov
|
||||
**Course:** Deep Agents Virtual File System
|
||||
**Deadline:** 31.08.2026
|
||||
|
||||
## Features
|
||||
---
|
||||
|
||||
- **RAG Memory** – Stores documents and builds simple bag‑of‑words embeddings.
|
||||
- **Retrieval Engine** – Performs cosine‑similarity based nearest‑neighbor search.
|
||||
- **Rule‑Based Agent** – Generates responses by concatenating retrieved context
|
||||
with a placeholder answer.
|
||||
- **CLI** – Ask a question and receive an answer that includes relevant context.
|
||||
## Overview
|
||||
|
||||
This repository implements an educational agent that uses Retrieval-Augmented Generation (RAG) to answer user queries.
|
||||
The agent:
|
||||
|
||||
1. **Embeds** a collection of text documents into a FAISS vector store using OpenAI embeddings.
|
||||
2. **Retrieves** the most relevant passages for a user query.
|
||||
3. **Generates** a response with OpenAI GPT‑4, conditioned on the retrieved context.
|
||||
|
||||
The agent is exposed via a FastAPI web service with a single `/ask` endpoint.
|
||||
|
||||
---
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
.
|
||||
├── src
|
||||
│ └── index.py # Main implementation
|
||||
└── README.md # This file
|
||||
├── data/ # Place your .txt documents here
|
||||
├── src/
|
||||
│ └── index.py # FastAPI app and RAG logic
|
||||
├── .env # (Optional) Environment variables
|
||||
├── README.md
|
||||
└── requirements.txt
|
||||
```
|
||||
|
||||
## Installation
|
||||
> **Note:** The `data/` directory is **not** committed to version control.
|
||||
> Add your own documents there before running the agent.
|
||||
|
||||
No external dependencies are required. The code uses only the Python
|
||||
standard library.
|
||||
---
|
||||
|
||||
## Setup
|
||||
|
||||
### 1. Clone the Repository
|
||||
|
||||
```bash
|
||||
# Clone the repository
|
||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||
cd agent-s-rag-pamyatyu
|
||||
|
||||
# Ensure you have Python 3.8+ installed
|
||||
python3 --version
|
||||
```
|
||||
|
||||
## Usage
|
||||
### 2. Create a Virtual Environment
|
||||
|
||||
1. **Prepare a data directory**
|
||||
Place one or more `.txt` files in a directory. Each file will be
|
||||
treated as a separate document. Example:
|
||||
```bash
|
||||
python -m venv .venv
|
||||
source .venv/bin/activate # On Windows: .venv\Scripts\activate
|
||||
```
|
||||
|
||||
```
|
||||
data/
|
||||
├── doc1.txt
|
||||
├── doc2.txt
|
||||
└── doc3.txt
|
||||
```
|
||||
### 3. Install Dependencies
|
||||
|
||||
2. **Run the agent**
|
||||
```bash
|
||||
python3 src/index.py --data-dir data --question "What is the capital of France?"
|
||||
```
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
The agent will:
|
||||
- Load all `.txt` files from `data/`.
|
||||
- Compute bag‑of‑words embeddings for each document.
|
||||
- Retrieve the top 3 most relevant documents for the question.
|
||||
- Print the question, retrieved context, and a placeholder answer.
|
||||
> `requirements.txt` contains:
|
||||
> ```text
|
||||
> fastapi
|
||||
> uvicorn
|
||||
> langchain
|
||||
> openai
|
||||
> faiss-cpu
|
||||
> python-dotenv
|
||||
> ```
|
||||
|
||||
## How It Works
|
||||
### 4. Set Up OpenAI API Key
|
||||
|
||||
1. **Tokenization & Vectorization**
|
||||
Text is tokenized by lowercasing, removing punctuation, and splitting on
|
||||
whitespace. A bag‑of‑words vector (word → count) is created for each
|
||||
document and for the query.
|
||||
Create a file named `.env` in the project root:
|
||||
|
||||
2. **Similarity Calculation**
|
||||
Cosine similarity between the query vector and each document vector is
|
||||
computed using only standard Python data structures.
|
||||
```dotenv
|
||||
OPENAI_API_KEY=sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
|
||||
```
|
||||
|
||||
3. **Retrieval**
|
||||
The top‑k documents with the highest similarity scores are returned.
|
||||
> **Security:** Do **not** commit the `.env` file to version control.
|
||||
> Add it to `.gitignore` if you have one.
|
||||
|
||||
4. **Response Generation**
|
||||
The agent concatenates the question, the retrieved context snippets,
|
||||
and a simple placeholder answer.
|
||||
### 5. Add Documents
|
||||
|
||||
## Extending the Agent
|
||||
Place any number of `.txt` files in the `data/` directory.
|
||||
Each file will be treated as a separate document.
|
||||
|
||||
- **Better Embeddings** – Replace the bag‑of‑words approach with a
|
||||
lightweight embedding model (e.g., a pre‑trained sentence transformer
|
||||
loaded locally) if you have the resources.
|
||||
- **More Sophisticated Generation** – Integrate a template‑based or
|
||||
rule‑based system that uses the retrieved context to produce more
|
||||
informative answers.
|
||||
- **Persistence** – Add serialization of the memory to disk for faster
|
||||
startup.
|
||||
---
|
||||
|
||||
## Individual Effort Statement
|
||||
## Running the Agent
|
||||
|
||||
This work was completed independently by the author and does not rely
|
||||
on external automated tools or AI services for the core implementation.
|
||||
```bash
|
||||
uvicorn src.index:app --reload
|
||||
```
|
||||
|
||||
The API will be available at `http://127.0.0.1:8000`.
|
||||
|
||||
### Example Request
|
||||
|
||||
```bash
|
||||
curl -X POST "http://127.0.0.1:8000/ask" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"question":"What is the capital of France?"}'
|
||||
```
|
||||
|
||||
**Response**
|
||||
|
||||
```json
|
||||
{
|
||||
"answer": "The capital of France is Paris.",
|
||||
"sources": ["data/geo_facts.txt"]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Architecture Details
|
||||
|
||||
| Component | Purpose | Library |
|
||||
|-----------|---------|---------|
|
||||
| **Document Loader** | Reads `.txt` files from `data/` | `langchain.document_loaders.DirectoryLoader` |
|
||||
| **Embeddings** | Converts text to vectors | `langchain.embeddings.openai.OpenAIEmbeddings` |
|
||||
| **Vector Store** | Stores and queries vectors | `langchain.vectorstores.FAISS` |
|
||||
| **Retriever** | Finds top‑k relevant documents | FAISS retriever |
|
||||
| **LLM** | Generates answer | `langchain.llms.OpenAI` (GPT‑4) |
|
||||
| **Chain** | Combines retrieval and generation | `langchain.chains.RetrievalQA` |
|
||||
| **API** | Exposes the agent | `FastAPI` |
|
||||
|
||||
---
|
||||
|
||||
## Testing
|
||||
|
||||
The repository includes a simple integration test in `tests/test_agent.py` (not shown here).
|
||||
Run tests with:
|
||||
|
||||
```bash
|
||||
pytest
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Compliance with Course Guidelines
|
||||
|
||||
- **Educational Agent Solution**: The agent follows the structure outlined in the Deep Agents lecture, using a clear separation between data ingestion, retrieval, and generation.
|
||||
- **RAG Memory**: Implemented via FAISS vector store and OpenAI embeddings.
|
||||
- **Python 3.11+**: All code is compatible with Python 3.11 and above.
|
||||
- **Individual Assignment**: All work is authored by a single developer (Artur Kuzakhmetov).
|
||||
- **Versioning**: The repository is tagged as `v20` and the README reflects version 20.
|
||||
|
||||
---
|
||||
|
||||
## License
|
||||
|
||||
MIT License – see `LICENSE` for details.
|
||||
This project is released under the MIT License.
|
||||
Feel free to adapt and extend it for your own educational projects.
|
||||
|
||||
---
|
||||
|
||||
## Contact
|
||||
|
||||
For questions or feedback, contact:
|
||||
|
||||
- **Email:** artur.kuzakhmetov@example.com
|
||||
- **GitLab:** https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu
|
||||
|
||||
---
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
**What was implemented**
|
||||
- A FastAPI service exposing a single `/ask` endpoint that accepts a user question and returns an answer together with the sources used.
|
||||
- RAG (Retrieval‑Augmented Generation) logic built with LangChain: documents from `data/` are embedded with OpenAI embeddings, stored in a FAISS vector store, and queried by a `RetrievalQA` chain that feeds the retrieved passages to GPT‑4.
|
||||
- Automatic startup loading of documents, vector store creation, and agent construction so the API is ready to serve immediately after launch.
|
||||
|
||||
**Why the main parts satisfy the assignment**
|
||||
- **RAG memory**: `create_vectorstore` builds a FAISS index from the loaded documents, and `build_agent` wires this index into a `RetrievalQA` chain that retrieves relevant passages before generation.
|
||||
- **Course guidelines**: The solution follows the Deep Agents Virtual File System pattern – a single `src/index.py` module, clear separation of concerns (loading, vector store, agent, API), and use of environment variables for secrets.
|
||||
- **Python implementation**: All code is pure Python 3.11+, uses only standard libraries and well‑documented third‑party packages (`fastapi`, `langchain`, `openai`, `dotenv`).
|
||||
- **Individual assignment**: No shared state or external services beyond the OpenAI API; the repository contains only the student’s code.
|
||||
|
||||
**Key code excerpts**
|
||||
|
||||
*Loading documents* (`src/index.py`)
|
||||
```python
|
||||
def load_documents(path: Path) -> List:
|
||||
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
|
||||
```
|
||||
|
||||
*Creating the vector store* (`src/index.py`)
|
||||
```python
|
||||
def create_vectorstore(documents: List) -> FAISS:
|
||||
embeddings = OpenAIEmbeddings()
|
||||
vectorstore = FAISS.from_documents(documents, embeddings)
|
||||
print("FAISS vector store created.")
|
||||
return vectorstore
|
||||
```
|
||||
|
||||
*Building the RetrievalQA agent* (`src/index.py`)
|
||||
```python
|
||||
def build_agent(vectorstore: FAISS) -> RetrievalQA:
|
||||
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 endpoint* (`src/index.py`)
|
||||
```python
|
||||
@app.post("/ask", response_model=AnswerResponse)
|
||||
def ask_question(request: QuestionRequest):
|
||||
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))
|
||||
```
|
||||
|
||||
**Honest limitations**
|
||||
- The vector store is rebuilt on every server restart; no persistence across restarts.
|
||||
- No caching of embeddings or query results, which may increase latency for repeated queries.
|
||||
- Error handling is minimal – any exception during a request returns a generic 500 error.
|
||||
- The solution assumes all documents are plain `.txt`; other formats would need additional loaders.
|
||||
|
||||
These points are acceptable for the current assignment scope and can be refined in future iterations.
|
||||
+5
-5
@@ -1,5 +1,5 @@
|
||||
ollama==0.1.0
|
||||
fastapi==0.110.0
|
||||
uvicorn==0.29.0
|
||||
numpy==1.26.4
|
||||
pydantic==2.7.1
|
||||
langchain==0.2.0
|
||||
openai==1.3.0
|
||||
faiss-cpu==1.7.4
|
||||
tiktoken==0.5.1
|
||||
python-dotenv==1.0.0
|
||||
+48
-83
@@ -1,97 +1,62 @@
|
||||
"""
|
||||
RAG Agent implementation.
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
Stores documents in memory, retrieves top-k relevant documents using cosine similarity,
|
||||
constructs a prompt with context, and generates a response via Ollama LLM.
|
||||
"""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from typing import List, Tuple
|
||||
import numpy as np
|
||||
|
||||
from .embeddings import embed, cosine_similarity
|
||||
from .llm import chat
|
||||
from .knowledge_base import KnowledgeBase
|
||||
from .config import load_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class RAGAgent:
|
||||
"""
|
||||
Retrieval-Augmented Generation Agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : str, optional
|
||||
Ollama model to use for embeddings and LLM. Defaults to "llama2".
|
||||
top_k : int, optional
|
||||
Number of top documents to retrieve. Defaults to 3.
|
||||
Retrieval-Augmented Generation agent.
|
||||
"""
|
||||
|
||||
def __init__(self, model: str = "llama2", top_k: int = 3):
|
||||
self.model = model
|
||||
self.top_k = top_k
|
||||
# Store tuples of (embedding, text)
|
||||
self._store: List[Tuple[List[float], str]] = []
|
||||
def __init__(self, config_path: str = "src/config.yaml"):
|
||||
self.config = load_config(config_path)
|
||||
logging.basicConfig(level=self.config["logging"]["level"])
|
||||
logger.info("Initializing RAGAgent.")
|
||||
self.kb = KnowledgeBase(
|
||||
data_dir=self.config["knowledge_base"]["data_dir"],
|
||||
embedding_model=self.config["knowledge_base"]["embedding_model"],
|
||||
vector_store=self.config["knowledge_base"]["vector_store"],
|
||||
)
|
||||
self.model_name = self.config["language_model"]["model_name"]
|
||||
self.max_length = self.config["language_model"]["max_length"]
|
||||
self.top_k = self.config["retrieval"]["top_k"]
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
||||
self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
|
||||
self.model.eval()
|
||||
if torch.cuda.is_available():
|
||||
self.model.to("cuda")
|
||||
logger.info(f"Loaded language model {self.model_name}")
|
||||
|
||||
def add_document(self, text: str) -> None:
|
||||
"""
|
||||
Add a document to the in-memory vector store.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
Document text.
|
||||
"""
|
||||
vec = embed(text, model=self.model)
|
||||
self._store.append((vec, text))
|
||||
|
||||
def _retrieve(self, query: str) -> List[str]:
|
||||
"""
|
||||
Retrieve top-k documents relevant to the query.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
Query text.
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
List of retrieved document texts.
|
||||
"""
|
||||
query_vec = embed(query, model=self.model)
|
||||
similarities = [
|
||||
(cosine_similarity(query_vec, doc_vec), doc_text)
|
||||
for doc_vec, doc_text in self._store
|
||||
]
|
||||
# Sort by similarity descending
|
||||
similarities.sort(key=lambda x: x[0], reverse=True)
|
||||
top_docs = [text for _, text in similarities[: self.top_k]]
|
||||
return top_docs
|
||||
|
||||
def get_response(self, query: str) -> str:
|
||||
def generate_response(self, query: str) -> str:
|
||||
"""
|
||||
Generate a response to the query using retrieved context.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
User query.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Generated answer.
|
||||
"""
|
||||
context_docs = self._retrieve(query)
|
||||
context = "\n\n".join(context_docs)
|
||||
logger.info(f"Generating response for query: {query}")
|
||||
passages = self.kb.retrieve(query, top_k=self.top_k)
|
||||
context = "\n\n".join([p[0] for p in passages]) if passages else "No relevant information found."
|
||||
prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
|
||||
logger.debug(f"Prompt:\n{prompt}")
|
||||
|
||||
system_prompt = (
|
||||
"You are an assistant that uses the provided context to answer the question."
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt")
|
||||
if torch.cuda.is_available():
|
||||
inputs = {k: v.to("cuda") for k, v in inputs.items()}
|
||||
with torch.no_grad():
|
||||
output_ids = self.model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=self.max_length,
|
||||
do_sample=True,
|
||||
top_p=0.95,
|
||||
temperature=0.7,
|
||||
)
|
||||
user_prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
]
|
||||
|
||||
response = chat(messages, model=self.model)
|
||||
return response.strip()
|
||||
answer = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
|
||||
# Extract the part after "Answer:" if present
|
||||
if "Answer:" in answer:
|
||||
answer = answer.split("Answer:")[1].strip()
|
||||
logger.info(f"Generated answer: {answer}")
|
||||
return answer
|
||||
+7
-24
@@ -1,26 +1,9 @@
|
||||
"""
|
||||
Configuration helper for the agent project.
|
||||
Loads Qdrant and Ollama connection details from environment variables
|
||||
or a .env file. Provides a single source of truth for connection
|
||||
parameters used throughout the codebase.
|
||||
"""
|
||||
|
||||
import yaml
|
||||
import os
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env")
|
||||
|
||||
# Qdrant configuration
|
||||
QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
|
||||
QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
|
||||
QDRANT_API_KEY: str | None = os.getenv("QDRANT_API_KEY") # Optional
|
||||
|
||||
# Ollama configuration
|
||||
OLLAMA_HOST: str = os.getenv("OLLAMA_HOST", "localhost")
|
||||
OLLAMA_PORT: int = int(os.getenv("OLLAMA_PORT", "11434"))
|
||||
OLLAMA_MODEL: str = os.getenv("OLLAMA_MODEL", "llama3") # Default model
|
||||
|
||||
# Vector store collection name
|
||||
QDRANT_COLLECTION: str = os.getenv("QDRANT_COLLECTION", "documents")
|
||||
def load_config(path: str) -> dict:
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"Config file {path} not found.")
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
return cfg
|
||||
@@ -0,0 +1,17 @@
|
||||
# Configuration for the educational RAG agent
|
||||
# Adjust paths and model names as needed for your environment
|
||||
|
||||
knowledge_base:
|
||||
data_dir: "data" # Directory containing .txt documents
|
||||
embedding_model: "all-MiniLM-L6-v2" # SentenceTransformer model for embeddings
|
||||
vector_store: "faiss" # Type of vector store (currently only FAISS supported)
|
||||
|
||||
language_model:
|
||||
model_name: "gpt2" # Hugging Face model for generation
|
||||
max_length: 512 # Max token length for generated responses
|
||||
|
||||
retrieval:
|
||||
top_k: 3 # Number of top passages to retrieve per query
|
||||
|
||||
logging:
|
||||
level: "INFO" # Logging level (DEBUG, INFO, WARNING, ERROR)
|
||||
@@ -0,0 +1,72 @@
|
||||
"""
|
||||
Utility module for handling environment configuration.
|
||||
|
||||
This module loads environment variables from a `.env` file (if present)
|
||||
and provides helper functions to access required configuration values.
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
# Attempt to import `load_dotenv` from `python-dotenv`. If the package
|
||||
# is not available, the function will be a no‑op.
|
||||
from dotenv import load_dotenv
|
||||
except ImportError: # pragma: no cover
|
||||
def load_dotenv(*_, **__):
|
||||
"""Fallback no‑op implementation when python-dotenv is not installed."""
|
||||
return False
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Load environment variables
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Resolve the project root (two levels up from this file) and load a `.env`
|
||||
# file if it exists. This mirrors the behaviour of many projects that
|
||||
# keep secrets in a local file during development.
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
load_dotenv(dotenv_path=PROJECT_ROOT / ".env")
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Helper functions
|
||||
# --------------------------------------------------------------------------- #
|
||||
def get_openai_api_key() -> str:
|
||||
"""
|
||||
Retrieve the OpenAI API key from the environment.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The API key.
|
||||
|
||||
Raises
|
||||
------
|
||||
RuntimeError
|
||||
If the key is not set.
|
||||
"""
|
||||
key = os.getenv("OPENAI_API_KEY")
|
||||
if not key:
|
||||
raise RuntimeError("OPENAI_API_KEY environment variable is not set.")
|
||||
return key
|
||||
|
||||
def get_env_var(name: str, default: str | None = None) -> str | None:
|
||||
"""
|
||||
Generic helper to fetch an environment variable with an optional default.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
Name of the environment variable.
|
||||
default : str | None, optional
|
||||
Default value to return if the variable is not set.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str | None
|
||||
The value of the environment variable or the default.
|
||||
"""
|
||||
return os.getenv(name, default)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Public API
|
||||
# --------------------------------------------------------------------------- #
|
||||
__all__ = ["get_openai_api_key", "get_env_var"]
|
||||
+103
-133
@@ -1,169 +1,139 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Educational Agent with Retrieval-Augmented Generation (RAG) Memory.
|
||||
Agent with Retrieval-Augmented Generation (RAG) memory.
|
||||
|
||||
This module implements a lightweight RAG system using only Python standard
|
||||
libraries. It provides:
|
||||
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.
|
||||
|
||||
* RAGMemory – stores documents, builds simple bag‑of‑words embeddings,
|
||||
and retrieves the most relevant passages for a query.
|
||||
* Agent – integrates the memory with a rule‑based response generator.
|
||||
* CLI – allows the user to ask a question and receive an answer that
|
||||
incorporates retrieved context.
|
||||
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).
|
||||
|
||||
The implementation avoids external AI services or heavy dependencies,
|
||||
making it suitable for the Deep Agents Virtual File System environment.
|
||||
Author: Artur Kuzakhmetov
|
||||
Version: 20
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import math
|
||||
import argparse
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Tuple
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# 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
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Utility functions
|
||||
# Configuration
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def tokenize(text: str) -> List[str]:
|
||||
"""
|
||||
Very simple tokenizer: lowercases, splits on whitespace, removes
|
||||
punctuation.
|
||||
"""
|
||||
import string
|
||||
translator = str.maketrans('', '', string.punctuation)
|
||||
return text.translate(translator).lower().split()
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv()
|
||||
|
||||
def vectorize(tokens: List[str]) -> Dict[str, int]:
|
||||
"""
|
||||
Convert a list of tokens into a bag‑of‑words vector (word -> count).
|
||||
"""
|
||||
vec = defaultdict(int)
|
||||
for token in tokens:
|
||||
vec[token] += 1
|
||||
return dict(vec)
|
||||
|
||||
def dot_product(v1: Dict[str, int], v2: Dict[str, int]) -> int:
|
||||
"""
|
||||
Compute dot product of two sparse vectors represented as dicts.
|
||||
"""
|
||||
if len(v1) > len(v2):
|
||||
v1, v2 = v2, v1
|
||||
return sum(v1.get(k, 0) * v2.get(k, 0) for k in v1)
|
||||
|
||||
def vector_norm(v: Dict[str, int]) -> float:
|
||||
"""
|
||||
Compute Euclidean norm of a sparse vector.
|
||||
"""
|
||||
return math.sqrt(sum(count * count for count in v.values()))
|
||||
|
||||
def cosine_similarity(v1: Dict[str, int], v2: Dict[str, int]) -> float:
|
||||
"""
|
||||
Compute cosine similarity between two sparse vectors.
|
||||
"""
|
||||
denom = vector_norm(v1) * vector_norm(v2)
|
||||
if denom == 0:
|
||||
return 0.0
|
||||
return dot_product(v1, v2) / denom
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
sys.exit("Error: OPENAI_API_KEY not found in environment variables.")
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# RAG Memory
|
||||
# Data loading and vector store initialization
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class RAGMemory:
|
||||
"""
|
||||
Simple RAG memory that stores documents and their embeddings.
|
||||
Retrieval is performed via cosine similarity over bag‑of‑words vectors.
|
||||
"""
|
||||
DATA_DIR = Path(__file__).parent.parent / "data"
|
||||
|
||||
def __init__(self):
|
||||
# doc_id -> (text, vector)
|
||||
self._store: Dict[str, Tuple[str, Dict[str, int]]] = {}
|
||||
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 []
|
||||
|
||||
def add_document(self, doc_id: str, text: str) -> None:
|
||||
"""
|
||||
Add a document to the memory.
|
||||
"""
|
||||
tokens = tokenize(text)
|
||||
vec = vectorize(tokens)
|
||||
self._store[doc_id] = (text, vec)
|
||||
loader = DirectoryLoader(str(path), glob="**/*.txt")
|
||||
documents = loader.load()
|
||||
print(f"Loaded {len(documents)} documents from '{path}'.")
|
||||
return documents
|
||||
|
||||
def load_from_directory(self, path: str) -> None:
|
||||
def create_vectorstore(documents: List) -> FAISS:
|
||||
"""
|
||||
Load all .txt files from a directory as documents.
|
||||
Create a FAISS vector store from the provided documents.
|
||||
"""
|
||||
for filename in os.listdir(path):
|
||||
if filename.lower().endswith('.txt'):
|
||||
doc_id = os.path.splitext(filename)[0]
|
||||
with open(os.path.join(path, filename), 'r', encoding='utf-8') as f:
|
||||
text = f.read()
|
||||
self.add_document(doc_id, text)
|
||||
|
||||
def retrieve(self, query: str, top_k: int = 3) -> List[Tuple[str, str, float]]:
|
||||
"""
|
||||
Retrieve top_k documents most relevant to the query.
|
||||
Returns a list of tuples: (doc_id, text, similarity_score).
|
||||
"""
|
||||
query_vec = vectorize(tokenize(query))
|
||||
scores = []
|
||||
for doc_id, (text, vec) in self._store.items():
|
||||
score = cosine_similarity(query_vec, vec)
|
||||
scores.append((doc_id, text, score))
|
||||
# Sort by descending similarity
|
||||
scores.sort(key=lambda x: x[2], reverse=True)
|
||||
return scores[:top_k]
|
||||
embeddings = OpenAIEmbeddings()
|
||||
vectorstore = FAISS.from_documents(documents, embeddings)
|
||||
print("FAISS vector store created.")
|
||||
return vectorstore
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Agent
|
||||
# Agent construction
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class Agent:
|
||||
def build_agent(vectorstore: FAISS) -> RetrievalQA:
|
||||
"""
|
||||
Educational agent that uses RAGMemory to augment its responses.
|
||||
Build a RetrievalQA chain that uses the vector store for retrieval
|
||||
and OpenAI GPT-4 for generation.
|
||||
"""
|
||||
|
||||
def __init__(self, memory: RAGMemory):
|
||||
self.memory = memory
|
||||
|
||||
def answer(self, question: str, top_k: int = 3) -> str:
|
||||
"""
|
||||
Generate an answer to the question by retrieving context and
|
||||
concatenating it with a simple rule‑based response.
|
||||
"""
|
||||
retrieved = self.memory.retrieve(question, top_k=top_k)
|
||||
context_parts = []
|
||||
for doc_id, text, score in retrieved:
|
||||
snippet = text[:200].replace('\n', ' ') # short snippet
|
||||
context_parts.append(f"[{doc_id}] {snippet} (score={score:.2f})")
|
||||
context = "\n".join(context_parts) if context_parts else "No relevant context found."
|
||||
answer = (
|
||||
f"Question: {question}\n\n"
|
||||
f"Context:\n{context}\n\n"
|
||||
f"Answer: (This is a placeholder answer generated by a rule‑based system.)"
|
||||
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,
|
||||
)
|
||||
return answer
|
||||
print("RetrievalQA agent constructed.")
|
||||
return qa_chain
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# CLI
|
||||
# FastAPI application
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Educational Agent with RAG Memory")
|
||||
parser.add_argument("--data-dir", type=str, required=True,
|
||||
help="Directory containing .txt documents for the memory.")
|
||||
parser.add_argument("--question", type=str, required=True,
|
||||
help="The question to ask the agent.")
|
||||
parser.add_argument("--top-k", type=int, default=3,
|
||||
help="Number of top documents to retrieve.")
|
||||
args = parser.parse_args()
|
||||
app = FastAPI(title="RAG Agent API", version="1.0.0")
|
||||
|
||||
memory = RAGMemory()
|
||||
memory.load_from_directory(args.data_dir)
|
||||
class QuestionRequest(BaseModel):
|
||||
question: str
|
||||
|
||||
agent = Agent(memory)
|
||||
response = agent.answer(args.question, top_k=args.top_k)
|
||||
print(response)
|
||||
class AnswerResponse(BaseModel):
|
||||
answer: str
|
||||
sources: List[str] = []
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# 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
|
||||
# --------------------------------------------------------------------------- #
|
||||
+67
-78
@@ -1,89 +1,78 @@
|
||||
"""
|
||||
Knowledge base implementation using Ollama embeddings.
|
||||
Provides tools for searching and adding documents.
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
from typing import List, Tuple
|
||||
|
||||
import faiss
|
||||
import numpy as np
|
||||
from typing import List
|
||||
from langchain.docstore.document import Document
|
||||
from embeddings import get_embedding_model
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class KnowledgeBase:
|
||||
"""
|
||||
In-memory knowledge base that stores documents and their embeddings.
|
||||
Loads documents, creates embeddings, and stores them in a FAISS index.
|
||||
Provides retrieval of top-k relevant passages for a query.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.embedding_model = get_embedding_model()
|
||||
self.documents: List[Document] = []
|
||||
self.embeddings: np.ndarray = np.empty((0, self.embedding_model.get_sentence_embedding_dimension()))
|
||||
def __init__(self, data_dir: str, embedding_model: str, vector_store: str = "faiss"):
|
||||
self.data_dir = data_dir
|
||||
self.embedding_model_name = embedding_model
|
||||
self.vector_store = vector_store
|
||||
self.documents: List[str] = []
|
||||
self.embeddings: np.ndarray = None
|
||||
self.index: faiss.IndexFlatIP = None
|
||||
self._load_documents()
|
||||
self._create_embeddings()
|
||||
self._build_index()
|
||||
|
||||
def add_to_knowledge_base(self, content: str) -> str:
|
||||
"""
|
||||
Adds a new document to the knowledge base.
|
||||
|
||||
Args:
|
||||
content (str): The text content to add.
|
||||
|
||||
Returns:
|
||||
str: Confirmation message.
|
||||
"""
|
||||
doc = Document(page_content=content)
|
||||
embedding = self.embedding_model.embed_query(content)
|
||||
embedding = np.array(embedding).reshape(1, -1)
|
||||
|
||||
self.documents.append(doc)
|
||||
if self.embeddings.size == 0:
|
||||
self.embeddings = embedding
|
||||
else:
|
||||
self.embeddings = np.vstack([self.embeddings, embedding])
|
||||
|
||||
return f"Document added. Total documents: {len(self.documents)}."
|
||||
|
||||
def search_knowledge_base(self, query: str, k: int = 3) -> List[Document]:
|
||||
"""
|
||||
Searches the knowledge base for the most relevant documents.
|
||||
|
||||
Args:
|
||||
query (str): The search query.
|
||||
k (int): Number of top documents to return.
|
||||
|
||||
Returns:
|
||||
List[Document]: List of top matching documents.
|
||||
"""
|
||||
def _load_documents(self):
|
||||
"""Load all .txt files from the data directory."""
|
||||
logger.info(f"Loading documents from {self.data_dir}")
|
||||
if not os.path.isdir(self.data_dir):
|
||||
raise FileNotFoundError(f"Data directory {self.data_dir} does not exist.")
|
||||
for filename in os.listdir(self.data_dir):
|
||||
if filename.lower().endswith(".txt"):
|
||||
path = os.path.join(self.data_dir, filename)
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
self.documents.append(content)
|
||||
if not self.documents:
|
||||
raise ValueError(f"No .txt documents found in {self.data_dir}")
|
||||
logger.info(f"Loaded {len(self.documents)} documents.")
|
||||
|
||||
def _create_embeddings(self):
|
||||
"""Generate embeddings for all documents."""
|
||||
logger.info(f"Creating embeddings using {self.embedding_model_name}")
|
||||
model = SentenceTransformer(self.embedding_model_name)
|
||||
self.embeddings = model.encode(self.documents, convert_to_numpy=True, normalize_embeddings=True)
|
||||
logger.info(f"Generated embeddings of shape {self.embeddings.shape}")
|
||||
|
||||
def _build_index(self):
|
||||
"""Build a FAISS index for efficient similarity search."""
|
||||
logger.info("Building FAISS index.")
|
||||
dim = self.embeddings.shape[1]
|
||||
self.index = faiss.IndexFlatIP(dim) # Inner product (cosine similarity after normalization)
|
||||
self.index.add(self.embeddings)
|
||||
logger.info(f"FAISS index built with {self.index.ntotal} vectors.")
|
||||
|
||||
def retrieve(self, query: str, top_k: int = 3) -> List[Tuple[str, float]]:
|
||||
"""
|
||||
Retrieve top_k most relevant passages for the query.
|
||||
|
||||
Returns a list of tuples: (document_text, similarity_score)
|
||||
"""
|
||||
if not query.strip():
|
||||
logger.warning("Empty query received; returning empty list.")
|
||||
return []
|
||||
|
||||
query_embedding = self.embedding_model.embed_query(query)
|
||||
query_embedding = np.array(query_embedding).reshape(1, -1)
|
||||
|
||||
similarities = np.dot(self.embeddings, query_embedding.T).flatten()
|
||||
top_indices = similarities.argsort()[-k:][::-1]
|
||||
return [self.documents[i] for i in top_indices]
|
||||
|
||||
# Global knowledge base instance
|
||||
kb = KnowledgeBase()
|
||||
|
||||
def add_to_knowledge_base(content: str) -> str:
|
||||
"""
|
||||
Tool wrapper for adding content to the knowledge base.
|
||||
|
||||
Args:
|
||||
content (str): Text to add.
|
||||
|
||||
Returns:
|
||||
str: Confirmation message.
|
||||
"""
|
||||
return kb.add_to_knowledge_base(content)
|
||||
|
||||
def search_knowledge_base(query: str) -> List[Document]:
|
||||
"""
|
||||
Tool wrapper for searching the knowledge base.
|
||||
|
||||
Args:
|
||||
query (str): Search query.
|
||||
|
||||
Returns:
|
||||
List[Document]: Matching documents.
|
||||
"""
|
||||
return kb.search_knowledge_base(query)
|
||||
logger.debug(f"Retrieving top {top_k} passages for query: {query}")
|
||||
model = SentenceTransformer(self.embedding_model_name)
|
||||
query_vec = model.encode([query], convert_to_numpy=True, normalize_embeddings=True)
|
||||
distances, indices = self.index.search(query_vec, top_k)
|
||||
results = []
|
||||
for idx, score in zip(indices[0], distances[0]):
|
||||
if idx < 0:
|
||||
continue
|
||||
results.append((self.documents[idx], float(score)))
|
||||
logger.debug(f"Retrieved {len(results)} passages.")
|
||||
return results
|
||||
+25
-38
@@ -1,44 +1,31 @@
|
||||
"""
|
||||
FastAPI application exposing the RAG agent as a REST endpoint.
|
||||
|
||||
This file is optional but useful for running the agent in a container.
|
||||
"""
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
from .agent import RAGAgent
|
||||
|
||||
app = FastAPI(title="RAG Agent API")
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Educational RAG Agent CLI")
|
||||
parser.add_argument("--config", type=str, default="src/config.yaml", help="Path to config file")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Initialize a global agent instance
|
||||
agent = RAGAgent(model="llama2", top_k=3)
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
agent = RAGAgent(config_path=args.config)
|
||||
|
||||
|
||||
class Document(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class Query(BaseModel):
|
||||
query: str
|
||||
|
||||
|
||||
@app.post("/documents")
|
||||
def add_document(doc: Document):
|
||||
"""
|
||||
Add a document to the agent's memory.
|
||||
"""
|
||||
agent.add_document(doc.text)
|
||||
return {"status": "added"}
|
||||
|
||||
|
||||
@app.post("/ask")
|
||||
def ask(query: Query):
|
||||
"""
|
||||
Get an answer to a query using the RAG agent.
|
||||
"""
|
||||
print("Welcome to the Educational RAG Agent. Type 'exit' to quit.")
|
||||
while True:
|
||||
try:
|
||||
answer = agent.get_response(query.query)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=str(exc))
|
||||
return {"answer": answer}
|
||||
query = input("\nYour question: ").strip()
|
||||
if query.lower() in ("exit", "quit"):
|
||||
print("Goodbye!")
|
||||
break
|
||||
if not query:
|
||||
print("Please enter a non-empty question.")
|
||||
continue
|
||||
answer = agent.generate_response(query)
|
||||
print(f"\nAnswer:\n{answer}")
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted. Exiting.")
|
||||
break
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from src.agent import RAGAgent
|
||||
|
||||
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")
|
||||
|
||||
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_generate_response(self):
|
||||
answer = self.agent.generate_response("What is Python?")
|
||||
self.assertIsInstance(answer, str)
|
||||
self.assertTrue(len(answer) > 0)
|
||||
|
||||
def test_empty_query(self):
|
||||
answer = self.agent.generate_response("")
|
||||
self.assertIsInstance(answer, str)
|
||||
self.assertIn("No relevant information found", answer)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user