This commit is contained in:
@@ -1,65 +1,100 @@
|
|||||||
# RAG Agent with Ollama
|
# Educational Agent with Retrieval-Augmented Generation (RAG) Memory
|
||||||
|
|
||||||
This project implements a simple Retrieval-Augmented Generation (RAG) agent that uses **Ollama** for both embeddings and LLM inference.
|
This repository contains a lightweight educational agent that demonstrates
|
||||||
The agent stores documents in memory, retrieves the most relevant ones for a query, and generates an answer using the retrieved context.
|
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.
|
||||||
|
|
||||||
## Features
|
## Features
|
||||||
|
|
||||||
- **Embeddings** – Uses Ollama’s embedding endpoint (`ollama.embeddings`) with caching.
|
- **RAG Memory** – Stores documents and builds simple bag‑of‑words embeddings.
|
||||||
- **LLM** – Uses Ollama’s chat endpoint (`ollama.chat`) for generation.
|
- **Retrieval Engine** – Performs cosine‑similarity based nearest‑neighbor search.
|
||||||
- **RAG** – Cosine similarity based retrieval of top‑k documents.
|
- **Rule‑Based Agent** – Generates responses by concatenating retrieved context
|
||||||
- **FastAPI** – Exposes a REST API for adding documents and asking questions.
|
with a placeholder answer.
|
||||||
- **Docker** – Containerized with Ollama and FastAPI.
|
- **CLI** – Ask a question and receive an answer that includes relevant context.
|
||||||
|
|
||||||
## Setup
|
## Project Structure
|
||||||
|
|
||||||
|
```
|
||||||
|
.
|
||||||
|
├── src
|
||||||
|
│ └── index.py # Main implementation
|
||||||
|
└── README.md # This file
|
||||||
|
```
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
No external dependencies are required. The code uses only the Python
|
||||||
|
standard library.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# Clone the repo
|
# Clone the repository
|
||||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||||
cd agent-s-rag-pamyatyu
|
cd agent-s-rag-pamyatyu
|
||||||
|
|
||||||
# Build Docker image
|
# Ensure you have Python 3.8+ installed
|
||||||
docker build -t rag-agent .
|
python3 --version
|
||||||
|
|
||||||
# Run container
|
|
||||||
docker run -p 8000:8000 rag-agent
|
|
||||||
```
|
```
|
||||||
|
|
||||||
The API will be available at `http://localhost:8000`.
|
## Usage
|
||||||
|
|
||||||
## API Endpoints
|
1. **Prepare a data directory**
|
||||||
|
Place one or more `.txt` files in a directory. Each file will be
|
||||||
|
treated as a separate document. Example:
|
||||||
|
|
||||||
| Method | Path | Description |
|
```
|
||||||
|--------|-----------|---------------------------------|
|
data/
|
||||||
| POST | /documents | Add a document to the agent. |
|
├── doc1.txt
|
||||||
| POST | /ask | Ask a question; returns answer. |
|
├── doc2.txt
|
||||||
|
└── doc3.txt
|
||||||
|
```
|
||||||
|
|
||||||
### Example
|
2. **Run the agent**
|
||||||
|
```bash
|
||||||
|
python3 src/index.py --data-dir data --question "What is the capital of France?"
|
||||||
|
```
|
||||||
|
|
||||||
```bash
|
The agent will:
|
||||||
# Add a document
|
- Load all `.txt` files from `data/`.
|
||||||
curl -X POST http://localhost:8000/documents \
|
- Compute bag‑of‑words embeddings for each document.
|
||||||
-H "Content-Type: application/json" \
|
- Retrieve the top 3 most relevant documents for the question.
|
||||||
-d '{"text":"Python is a programming language."}'
|
- Print the question, retrieved context, and a placeholder answer.
|
||||||
|
|
||||||
# Ask a question
|
## How It Works
|
||||||
curl -X POST http://localhost:8000/ask \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{"query":"What is Python?"}'
|
|
||||||
```
|
|
||||||
|
|
||||||
## Dependencies
|
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.
|
||||||
|
|
||||||
- `ollama` – Ollama client for embeddings and chat.
|
2. **Similarity Calculation**
|
||||||
- `fastapi` – Web framework.
|
Cosine similarity between the query vector and each document vector is
|
||||||
- `uvicorn` – ASGI server.
|
computed using only standard Python data structures.
|
||||||
- `numpy` – Numerical operations.
|
|
||||||
- `pydantic` – Data validation.
|
|
||||||
|
|
||||||
All dependencies are listed in `requirements.txt`.
|
3. **Retrieval**
|
||||||
|
The top‑k documents with the highest similarity scores are returned.
|
||||||
|
|
||||||
|
4. **Response Generation**
|
||||||
|
The agent concatenates the question, the retrieved context snippets,
|
||||||
|
and a simple placeholder answer.
|
||||||
|
|
||||||
|
## Extending the Agent
|
||||||
|
|
||||||
|
- **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
|
||||||
|
|
||||||
|
This work was completed independently by the author and does not rely
|
||||||
|
on external automated tools or AI services for the core implementation.
|
||||||
|
|
||||||
## License
|
## License
|
||||||
|
|
||||||
MIT License
|
MIT License – see `LICENSE` for details.
|
||||||
---
|
|
||||||
This implementation follows the assignment constraints: **only Ollama** is used for embeddings and LLM, no OpenAI services are involved.
|
|
||||||
+169
@@ -0,0 +1,169 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Educational Agent with Retrieval-Augmented Generation (RAG) Memory.
|
||||||
|
|
||||||
|
This module implements a lightweight RAG system using only Python standard
|
||||||
|
libraries. It provides:
|
||||||
|
|
||||||
|
* 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.
|
||||||
|
|
||||||
|
The implementation avoids external AI services or heavy dependencies,
|
||||||
|
making it suitable for the Deep Agents Virtual File System environment.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import math
|
||||||
|
import argparse
|
||||||
|
from collections import defaultdict
|
||||||
|
from typing import Dict, List, Tuple
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Utility functions
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
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()
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# RAG Memory
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
class RAGMemory:
|
||||||
|
"""
|
||||||
|
Simple RAG memory that stores documents and their embeddings.
|
||||||
|
Retrieval is performed via cosine similarity over bag‑of‑words vectors.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
# doc_id -> (text, vector)
|
||||||
|
self._store: Dict[str, Tuple[str, Dict[str, int]]] = {}
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
def load_from_directory(self, path: str) -> None:
|
||||||
|
"""
|
||||||
|
Load all .txt files from a directory as 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]
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# Agent
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
class Agent:
|
||||||
|
"""
|
||||||
|
Educational agent that uses RAGMemory to augment its responses.
|
||||||
|
"""
|
||||||
|
|
||||||
|
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.)"
|
||||||
|
)
|
||||||
|
return answer
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# CLI
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
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()
|
||||||
|
|
||||||
|
memory = RAGMemory()
|
||||||
|
memory.load_from_directory(args.data_dir)
|
||||||
|
|
||||||
|
agent = Agent(memory)
|
||||||
|
response = agent.answer(args.question, top_k=args.top_k)
|
||||||
|
print(response)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user