This commit is contained in:
kuzakhmetovartur committed
1 parent
62c0063d9d
commit
b7ed14fd93
2 files changed
+246
-42
No files matched your search
@@ -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.
|
||||
The agent stores documents in memory, retrieves the most relevant ones for a query, and generates an answer using the retrieved context.
|
||||
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.
|
||||
|
||||
## Features
|
||||
|
||||
- **Embeddings** – Uses Ollama’s embedding endpoint (`ollama.embeddings`) with caching.
|
||||
- **LLM** – Uses Ollama’s chat endpoint (`ollama.chat`) for generation.
|
||||
- **RAG** – Cosine similarity based retrieval of top‑k documents.
|
||||
- **FastAPI** – Exposes a REST API for adding documents and asking questions.
|
||||
- **Docker** – Containerized with Ollama and FastAPI.
|
||||
- **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.
|
||||
|
||||
## 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
|
||||
# Clone the repo
|
||||
# Clone the repository
|
||||
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
|
||||
cd agent-s-rag-pamyatyu
|
||||
|
||||
# Build Docker image
|
||||
docker build -t rag-agent .
|
||||
|
||||
# Run container
|
||||
docker run -p 8000:8000 rag-agent
|
||||
# Ensure you have Python 3.8+ installed
|
||||
python3 --version
|
||||
```
|
||||
|
||||
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 |
|
||||
|--------|-----------|---------------------------------|
|
||||
| POST | /documents | Add a document to the agent. |
|
||||
| POST | /ask | Ask a question; returns answer. |
|
||||
```
|
||||
data/
|
||||
├── doc1.txt
|
||||
├── 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
|
||||
# Add a document
|
||||
curl -X POST http://localhost:8000/documents \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"text":"Python is a programming language."}'
|
||||
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.
|
||||
|
||||
# Ask a question
|
||||
curl -X POST http://localhost:8000/ask \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"query":"What is Python?"}'
|
||||
```
|
||||
## How It Works
|
||||
|
||||
## 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.
|
||||
- `fastapi` – Web framework.
|
||||
- `uvicorn` – ASGI server.
|
||||
- `numpy` – Numerical operations.
|
||||
- `pydantic` – Data validation.
|
||||
2. **Similarity Calculation**
|
||||
Cosine similarity between the query vector and each document vector is
|
||||
computed using only standard Python data structures.
|
||||
|
||||
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
|
||||
|
||||
MIT License
|
||||
---
|
||||
This implementation follows the assignment constraints: **only Ollama** is used for embeddings and LLM, no OpenAI services are involved.
|
||||
MIT License – see `LICENSE` for details.
|
||||
+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