feat: solution for 'Практическое задание: Агент с RAG-памятью'
This commit is contained in:
@@ -0,0 +1,5 @@
|
|||||||
|
node_modules/
|
||||||
|
.env
|
||||||
|
dist/
|
||||||
|
build/
|
||||||
|
*.log
|
||||||
@@ -0,0 +1,93 @@
|
|||||||
|
```markdown
|
||||||
|
# RAG Agent with Qdrant and Ollama
|
||||||
|
|
||||||
|
This project implements an AI agent that can search and add documents to a local knowledge base using **Qdrant** for vector storage and **Ollama** for embeddings and LLM inference. The agent is built with **LangChain** and exposes two tools:
|
||||||
|
|
||||||
|
- `search_knowledge_base(query, max_results)` – semantic search in the knowledge base.
|
||||||
|
- `add_to_knowledge_base(content, title)` – add a new document to the knowledge base.
|
||||||
|
|
||||||
|
## Features
|
||||||
|
|
||||||
|
- **Vector store**: Qdrant with Ollama embeddings (`nomic-embed-text`).
|
||||||
|
- **Chunking**: Recursive character splitter with overlap.
|
||||||
|
- **Agent**: Zero-shot React agent that uses the two tools.
|
||||||
|
- **CLI**: Interactive command line interface to add documents and query the agent.
|
||||||
|
- **Batch loading**: Script to load all text files from a directory into the knowledge base.
|
||||||
|
|
||||||
|
## Prerequisites
|
||||||
|
|
||||||
|
- Python 3.10+
|
||||||
|
- Docker (for Qdrant) or a running Qdrant instance.
|
||||||
|
- Ollama installed locally with the following models:
|
||||||
|
```bash
|
||||||
|
ollama pull llama3
|
||||||
|
ollama pull nomic-embed-text
|
||||||
|
```
|
||||||
|
|
||||||
|
## Setup
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Clone the repository
|
||||||
|
git clone https://github.com/your-username/rag-agent.git
|
||||||
|
cd rag-agent
|
||||||
|
|
||||||
|
# Create a virtual environment
|
||||||
|
python -m venv venv
|
||||||
|
source venv/bin/activate # On Windows: venv\Scripts\activate
|
||||||
|
|
||||||
|
# Install dependencies
|
||||||
|
pip install -r requirements.txt
|
||||||
|
|
||||||
|
# Start Qdrant (Docker example)
|
||||||
|
docker run -p 6333:6333 qdrant/qdrant
|
||||||
|
```
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
### 1. Load documents into the knowledge base
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python src/main.py /path/to/documents
|
||||||
|
```
|
||||||
|
|
||||||
|
Supported file types: `.txt`, `.md`. (PDF support can be added with an additional parser.)
|
||||||
|
|
||||||
|
### 2. Start the interactive CLI
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python src/cli.py
|
||||||
|
```
|
||||||
|
|
||||||
|
Commands:
|
||||||
|
|
||||||
|
- `/add <file_path>` – Add a single document.
|
||||||
|
- `/search <query>` – Query the agent.
|
||||||
|
- `/quit` – Exit.
|
||||||
|
|
||||||
|
### 3. Example
|
||||||
|
|
||||||
|
```bash
|
||||||
|
> /add example.txt
|
||||||
|
Document 'example' added to knowledge base with 3 chunks.
|
||||||
|
> /search What is the capital of France?
|
||||||
|
1. The capital of France is Paris. (Title: example)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Project Structure
|
||||||
|
|
||||||
|
```
|
||||||
|
rag-agent/
|
||||||
|
├── src/
|
||||||
|
│ ├── agent.py
|
||||||
|
│ ├── cli.py
|
||||||
|
│ ├── main.py
|
||||||
|
│ ├── tools.py
|
||||||
|
│ └── vector_store.py
|
||||||
|
├── requirements.txt
|
||||||
|
└── README.md
|
||||||
|
```
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
MIT License
|
||||||
|
```
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
```
|
||||||
|
langchain
|
||||||
|
langchain-qdrant
|
||||||
|
langchain-ollama
|
||||||
|
qdrant-client
|
||||||
|
python-dotenv
|
||||||
|
```
|
||||||
@@ -0,0 +1,40 @@
|
|||||||
|
```python
|
||||||
|
"""
|
||||||
|
Agent creation with RAG integration.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from langchain.llms import Ollama
|
||||||
|
from langchain.agents import initialize_agent, AgentType
|
||||||
|
from langchain.tools import Tool
|
||||||
|
|
||||||
|
from tools import search_knowledge_base, add_to_knowledge_base
|
||||||
|
|
||||||
|
|
||||||
|
def create_agent():
|
||||||
|
"""
|
||||||
|
Create and configure the LangChain agent.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
AgentExecutor instance ready to run queries.
|
||||||
|
"""
|
||||||
|
llm = Ollama(model="llama3")
|
||||||
|
|
||||||
|
tools = [
|
||||||
|
Tool.from_function(search_knowledge_base),
|
||||||
|
Tool.from_function(add_to_knowledge_base),
|
||||||
|
]
|
||||||
|
|
||||||
|
agent = initialize_agent(
|
||||||
|
tools,
|
||||||
|
llm,
|
||||||
|
agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
|
||||||
|
verbose=True,
|
||||||
|
agent_kwargs={
|
||||||
|
"system_message": (
|
||||||
|
"You are an AI assistant that can search and add documents to a knowledge base. "
|
||||||
|
"Use the provided tools to answer user queries."
|
||||||
|
)
|
||||||
|
},
|
||||||
|
)
|
||||||
|
return agent
|
||||||
|
```
|
||||||
+62
@@ -0,0 +1,62 @@
|
|||||||
|
```python
|
||||||
|
"""
|
||||||
|
Interactive CLI for the RAG agent.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from tools import add_to_knowledge_base
|
||||||
|
from agent import create_agent
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
agent = create_agent()
|
||||||
|
print("RAG Agent CLI. Commands:")
|
||||||
|
print(" /add <file_path> - Add a document to the knowledge base.")
|
||||||
|
print(" /search <query> - Search the knowledge base.")
|
||||||
|
print(" /quit - Exit the program.")
|
||||||
|
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
inp = input("> ").strip()
|
||||||
|
except EOFError:
|
||||||
|
break
|
||||||
|
|
||||||
|
if not inp:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if inp.startswith("/add"):
|
||||||
|
parts = inp.split(maxsplit=1)
|
||||||
|
if len(parts) < 2:
|
||||||
|
print("Usage: /add <file_path>")
|
||||||
|
continue
|
||||||
|
file_path = Path(parts[1])
|
||||||
|
if not file_path.is_file():
|
||||||
|
print(f"File {file_path} does not exist.")
|
||||||
|
continue
|
||||||
|
content = file_path.read_text(encoding="utf-8")
|
||||||
|
title = file_path.stem
|
||||||
|
result = add_to_knowledge_base(content, title)
|
||||||
|
print(result)
|
||||||
|
|
||||||
|
elif inp.startswith("/search"):
|
||||||
|
parts = inp.split(maxsplit=1)
|
||||||
|
if len(parts) < 2:
|
||||||
|
print("Usage: /search <query>")
|
||||||
|
continue
|
||||||
|
query = parts[1]
|
||||||
|
response = agent.run(query)
|
||||||
|
print(response)
|
||||||
|
|
||||||
|
elif inp.startswith("/quit"):
|
||||||
|
print("Goodbye!")
|
||||||
|
break
|
||||||
|
|
||||||
|
else:
|
||||||
|
print("Unknown command. Use /add, /search, /quit.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
```
|
||||||
+42
@@ -0,0 +1,42 @@
|
|||||||
|
```python
|
||||||
|
"""
|
||||||
|
Script to load documents from a directory into the vector store.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from vector_store import vector_store
|
||||||
|
from langchain_text_splitter import RecursiveCharacterTextSplitter
|
||||||
|
|
||||||
|
# Chunking configuration
|
||||||
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
|
|
||||||
|
|
||||||
|
def load_documents_from_dir(directory: str):
|
||||||
|
"""
|
||||||
|
Load all supported text files from the given directory into the knowledge base.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
directory: Path to the directory containing documents.
|
||||||
|
"""
|
||||||
|
dir_path = Path(directory)
|
||||||
|
for file_path in dir_path.rglob("*"):
|
||||||
|
if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md"}:
|
||||||
|
content = file_path.read_text(encoding="utf-8")
|
||||||
|
title = file_path.stem
|
||||||
|
chunks = splitter.split_text(content)
|
||||||
|
vector_store.add_documents(chunks, [title] * len(chunks))
|
||||||
|
print(f"Loaded {file_path} into knowledge base.")
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="Load documents into the knowledge base.")
|
||||||
|
parser.add_argument("directory", help="Path to directory with documents.")
|
||||||
|
args = parser.parse_args()
|
||||||
|
load_documents_from_dir(args.directory)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
```
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
```python
|
||||||
|
"""
|
||||||
|
Tools for the RAG agent: searching and adding to the knowledge base.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from langchain.tools import tool
|
||||||
|
from langchain_text_splitter import RecursiveCharacterTextSplitter
|
||||||
|
|
||||||
|
from vector_store import vector_store
|
||||||
|
|
||||||
|
# Chunking configuration
|
||||||
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||||||
|
"""
|
||||||
|
Semantic search in the knowledge base.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query: Search query.
|
||||||
|
max_results: Maximum number of results to return.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Formatted string with search results.
|
||||||
|
"""
|
||||||
|
results = vector_store.search(query, max_results)
|
||||||
|
if not results:
|
||||||
|
return "No relevant documents found."
|
||||||
|
|
||||||
|
formatted = []
|
||||||
|
for i, doc in enumerate(results):
|
||||||
|
snippet = doc.page_content[:200].replace("\n", " ")
|
||||||
|
formatted.append(
|
||||||
|
f"{i + 1}. {snippet} (Title: {doc.metadata.get('title', 'N/A')})"
|
||||||
|
)
|
||||||
|
return "\n".join(formatted)
|
||||||
|
|
||||||
|
|
||||||
|
@tool
|
||||||
|
def add_to_knowledge_base(content: str, title: str) -> str:
|
||||||
|
"""
|
||||||
|
Add a document to the knowledge base.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content: Full text of the document.
|
||||||
|
title: Title of the document.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Confirmation message.
|
||||||
|
"""
|
||||||
|
chunks = splitter.split_text(content)
|
||||||
|
vector_store.add_documents(chunks, [title] * len(chunks))
|
||||||
|
return f"Document '{title}' added to knowledge base with {len(chunks)} chunks."
|
||||||
|
```
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
```python
|
||||||
|
"""
|
||||||
|
Vector store implementation using Qdrant and Ollama embeddings.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from langchain_ollama import OllamaEmbeddings
|
||||||
|
from langchain_qdrant import QdrantVectorStore
|
||||||
|
|
||||||
|
|
||||||
|
class QdrantVectorStoreWrapper:
|
||||||
|
"""
|
||||||
|
Wrapper around LangChain's QdrantVectorStore.
|
||||||
|
Handles initialization, document addition, and similarity search.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
collection_name: str = "knowledge_base",
|
||||||
|
host: str = "localhost",
|
||||||
|
port: int = 6333,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize the vector store.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
collection_name: Name of the Qdrant collection.
|
||||||
|
host: Qdrant host address.
|
||||||
|
port: Qdrant port.
|
||||||
|
"""
|
||||||
|
self.embedding = OllamaEmbeddings(model="nomic-embed-text")
|
||||||
|
self.store = QdrantVectorStore(
|
||||||
|
url=f"http://{host}:{port}",
|
||||||
|
collection_name=collection_name,
|
||||||
|
embedding=self.embedding,
|
||||||
|
)
|
||||||
|
|
||||||
|
def add_documents(self, documents: List[str], titles: List[str]) -> None:
|
||||||
|
"""
|
||||||
|
Add documents to the vector store with metadata.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
documents: List of document texts.
|
||||||
|
titles: List of titles corresponding to each document.
|
||||||
|
"""
|
||||||
|
metadatas = [{"title": title} for title in titles]
|
||||||
|
self.store.add_texts(documents, metadatas=metadatas)
|
||||||
|
|
||||||
|
def search(self, query: str, k: int = 5):
|
||||||
|
"""
|
||||||
|
Perform a similarity search.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query: Query string.
|
||||||
|
k: Number of results to return.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of Document objects sorted by relevance.
|
||||||
|
"""
|
||||||
|
return self.store.similarity_search(query, k)
|
||||||
|
|
||||||
|
|
||||||
|
# Global instance used by tools and agent
|
||||||
|
vector_store = QdrantVectorStoreWrapper()
|
||||||
|
```
|
||||||
Reference in New Issue
Block a user