kuzakhmetovartur/agent-s-rag-pamyatyu
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: solution for 'Агент с RAG-памятью'
CI / build (push) Has been cancelled
Details

Browse Source
This commit is contained in:
Артур Кузахметов
2026-06-30 15:36:47 +03:00
parent 1279deaaa6
commit 42fd00d924
4 changed files with 198 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+57 -65
View File
@@ -1,14 +1,16 @@
# Agent with RAG Memory
# RAG Agent with Ollama Embeddings
This project implements a retrieval‑augmented generation (RAG) agent that uses **Qdrant** as the vector store and **Ollama** for local LLM inference.
The agent follows the latest LangChain API and is fully configurable via environment variables.
This project demonstrates a simple Retrieval-Augmented Generation (RAG) agent that uses **OllamaEmbeddings** for vector similarity search and a local in‑memory knowledge base.
The agent is built with **LangChain** and exposes two tools:
## Features
- `search_knowledge_base`: Search the knowledge base for relevant documents.
- `add_to_knowledge_base`: Add new content to the knowledge base.
- **Qdrant** vector store (via `langchain-qdrant`)
- **Ollama** local LLM integration (via `langchain-ollama`)
- Retrieval‑augmented generation with conversation memory
- Simple CLI interface for quick testing
## Prerequisites
- Python 3.10+
- An Ollama server running locally (e.g., `ollama serve`).
- The Ollama model you want to use (default is `mistral`).
## Installation
@@ -17,82 +19,72 @@ The agent follows the latest LangChain API and is fully configurable via environ
git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
cd agent-s-rag-pamyatyu
# Create a virtual environment (optional but recommended)
python -m venv .venv
source .venv/bin/activate # On Windows: .venv\\Scripts\\activate
# Create a virtual environment
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
# Install dependencies
pip install -r requirements.txt
```
`requirements.txt` contains:
```text
langchain
langchain-community
openai
```
## Configuration
Create a `.env` file in the project root (or set environment variables directly):
```dotenv
# Qdrant
QDRANT_HOST=localhost
QDRANT_PORT=6333
QDRANT_API_KEY= # leave empty if no key
# Ollama
OLLAMA_HOST=localhost
OLLAMA_PORT=11434
OLLAMA_MODEL=llama3
# Optional: collection name
QDRANT_COLLECTION=documents
```
> **Note**: The Qdrant instance must be running and accessible at the specified host/port.
> The Ollama server must be running locally and expose the chosen model.
## Usage
### Adding Documents
```python
from src.agent import Agent
from langchain_core.documents import Document
agent = Agent()
docs = [
Document(page_content="Python is a programming language.", metadata={"source": "python.txt"}),
Document(page_content="LangChain is a framework for LLM applications.", metadata={"source": "langchain.txt"}),
]
agent.add_documents(docs)
```
### Querying the Agent
Set the Ollama model via environment variable (optional):
```bash
python -m src.agent "What is LangChain?"
export OLLAMA_MODEL=mistral # or any other model available in Ollama
```
or from Python:
If you run the Ollama server on a non‑default host/port, set:
```python
response = agent.run("What is LangChain?")
print(response["answer"])
```bash
export OLLAMA_HOST=http://localhost:11434
```
The response will include the answer and the source documents used.
## Running the Agent
## Project Structure
```bash
python src/agent.py
```
You will see a prompt:
```
agent-s-rag-pamyatyu/
├── src/
│ ├── agent.py
│ ├── config.py
│ └── vector_store.py
├── requirements.txt
├── pyproject.toml
└── README.md
Welcome to the RAG Agent. Type 'exit' to quit.
User:
```
- **Add knowledge**:
`add_to_knowledge_base This is a new piece of information.`
- **Search knowledge**:
`search_knowledge_base information`
The agent will automatically decide which tool to use based on the user query.
## Example Session
```
User: add_to_knowledge_base Python is a versatile programming language.
Agent: Document added. Total documents: 1.
User: search_knowledge_base programming language
Agent: Python is a versatile programming language.
```
## Notes
- The knowledge base is **in‑memory**; data will be lost when the program exits.
- For persistent storage, replace the in‑memory implementation with a vector database such as Chroma or FAISS.
- The LLM used for generation is OpenAI’s GPT‑3.5 via the `openai` package. Adjust the `OpenAI` initialization if you prefer another model.
## License
MIT License
src/agent.py
+59 -78
View File
@@ -1,94 +1,75 @@
"""
Agent implementation using the latest LangChain API.
The agent uses a RetrievalQA chain backed by the Qdrant vector store
and the Ollama LLM for local inference.
Agent implementation that uses LangChain to interact with the knowledge base.
"""
from __future__ import annotations
import os
from typing import List
from langchain import OpenAI
from langchain.agents import ZeroShotAgent, AgentExecutor
from langchain.tools import BaseTool
from knowledge_base import add_to_knowledge_base, search_knowledge_base
from langchain.schema import AgentAction, AgentFinish
from typing import Any
from langchain_ollama import Ollama
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from .config import OLLAMA_HOST, OLLAMA_PORT, OLLAMA_MODEL
from .vector_store import QdrantVectorStore
from langchain_core.documents import Document
class Agent:
class SearchTool(BaseTool):
"""
A simple retrieval-based agent that answers user queries
using documents stored in Qdrant and an Ollama LLM.
Tool to search the knowledge base.
"""
name = "search_knowledge_base"
description = (
"Search the knowledge base for relevant information. "
"Input: query string."
)
def __init__(self, collection_name: str = "documents"):
# Initialize LLM
self.llm = Ollama(
model=OLLAMA_MODEL,
base_url=f"http://{OLLAMA_HOST}:{OLLAMA_PORT}",
)
def _run(self, query: str) -> str:
docs: List = search_knowledge_base(query)
if not docs:
return "No relevant documents found."
return "\n---\n".join([doc.page_content for doc in docs])
# Initialize vector store
self.vector_store = QdrantVectorStore(collection_name=collection_name)
class AddTool(BaseTool):
"""
Tool to add new knowledge to the knowledge base.
"""
name = "add_to_knowledge_base"
description = (
"Add new knowledge to the knowledge base. "
"Input: content string."
)
# Memory to keep conversation context
self.memory = ConversationBufferMemory(memory_key="chat_history")
def _run(self, content: str) -> str:
return add_to_knowledge_base(content)
# Prompt template
self.prompt = PromptTemplate(
input_variables=["chat_history", "question"],
template=(
"You are a helpful assistant. Use the following context to answer the question.\n"
"Context:\n{chat_history}\n"
"Question: {question}\n"
"Answer:"
),
)
# Instantiate tools
tools = [SearchTool(), AddTool()]
# RetrievalQA chain
self.chain = RetrievalQA.from_chain_type(
llm=self.llm,
chain_type="stuff",
retriever=self.vector_store.as_retriever(),
memory=self.memory,
return_source_documents=True,
chain_type_kwargs={"prompt": self.prompt},
)
# LLM configuration
llm = OpenAI(temperature=0)
def add_documents(self, documents: list[Document]) -> None:
"""
Add documents to the underlying vector store.
"""
self.vector_store.add_documents(documents)
# Create the agent
agent = ZeroShotAgent(llm=llm, tools=tools)
def run(self, question: str) -> Any:
"""
Run the agent on a user question.
Returns the LLM's answer and the source documents.
"""
result = self.chain({"question": question})
return result
def __call__(self, question: str) -> Any:
return self.run(question)
# Executor that runs the agent
agent_executor = AgentExecutor.from_agent_and_tools(
agent=agent,
tools=tools,
verbose=True
)
def main() -> None:
"""
Simple REPL to interact with the agent.
"""
print("Welcome to the RAG Agent. Type 'exit' to quit.")
while True:
user_input = input("User: ")
if user_input.lower() in ("exit", "quit"):
print("Goodbye!")
break
try:
response = agent_executor.run(user_input)
print(f"Agent: {response}")
except Exception as e:
print(f"Error: {e}")
if __name__ == "__main__":
import json
import sys
# Simple CLI usage
agent = Agent()
if len(sys.argv) > 1:
query = " ".join(sys.argv[1:])
else:
query = input("Enter your question: ")
response = agent.run(query)
print("\nAnswer:\n", response["answer"])
print("\nSources:")
for doc in response["source_documents"]:
print(f"- {doc.metadata.get('source', 'unknown')}")
main()
src/embeddings.py
+18
View File
@@ -0,0 +1,18 @@
"""
Embeddings module that provides an OllamaEmbeddings instance.
"""
import os
from langchain_community.embeddings import OllamaEmbeddings
def get_embedding_model() -> OllamaEmbeddings:
"""
Returns an OllamaEmbeddings instance configured with the model name
specified by the OLLAMA_MODEL environment variable or defaults to
'mistral'.
Returns:
OllamaEmbeddings: The embedding model instance.
"""
model_name = os.getenv("OLLAMA_MODEL", "mistral")
return OllamaEmbeddings(model=model_name)
src/knowledge_base.py
+64 -106
View File
@@ -1,131 +1,89 @@
"""
Knowledge Base Tool for the Agent.
This module implements a simple file‑based knowledge base that can be
used by the agent and accessed via the CLI. The knowledge base is
stored as a JSON file (`knowledge_base.json`) in the same directory
as this module. Each entry is a key/value pair where the key is a
string and the value is any JSON‑serialisable object.
The class provides three public methods:
* add_entry(key, value) – Add or update an entry.
* query_entry(key) – Retrieve the value for a key.
* delete_entry(key) – Remove an entry.
The tool is intentionally lightweight and does not depend on any
external libraries beyond the Python standard library.
Knowledge base implementation using Ollama embeddings.
Provides tools for searching and adding documents.
"""
import json
import os
from pathlib import Path
from typing import Any, Dict, Optional
import numpy as np
from typing import List
from langchain.docstore.document import Document
from embeddings import get_embedding_model
class KnowledgeBaseTool:
class KnowledgeBase:
"""
A simple file‑based knowledge base tool.
In-memory knowledge base that stores documents and their embeddings.
"""
def __init__(self, storage_path: Optional[Path] = None) -> None:
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 add_to_knowledge_base(self, content: str) -> str:
"""
Initialise the knowledge base.
Adds a new document to the knowledge base.
Parameters
----------
storage_path : Optional[Path]
Path to the JSON file used for storage. If not provided,
a file named ``knowledge_base.json`` in the same directory
as this module is used.
Args:
content (str): The text content to add.
Returns:
str: Confirmation message.
"""
if storage_path is None:
storage_path = Path(__file__).parent / "knowledge_base.json"
self.storage_path = storage_path
# Ensure the storage file exists
if not self.storage_path.exists():
self.storage_path.write_text("{}")
doc = Document(page_content=content)
embedding = self.embedding_model.embed_query(content)
embedding = np.array(embedding).reshape(1, -1)
def _load(self) -> Dict[str, Any]:
"""Load the knowledge base from disk."""
try:
data = json.loads(self.storage_path.read_text())
if not isinstance(data, dict):
raise ValueError("Knowledge base file corrupted: not a dict")
return data
except json.JSONDecodeError:
raise ValueError("Knowledge base file corrupted: invalid JSON")
self.documents.append(doc)
if self.embeddings.size == 0:
self.embeddings = embedding
else:
self.embeddings = np.vstack([self.embeddings, embedding])
def _save(self, data: Dict[str, Any]) -> None:
"""Persist the knowledge base to disk."""
self.storage_path.write_text(json.dumps(data, indent=2, ensure_ascii=False))
return f"Document added. Total documents: {len(self.documents)}."
def add_entry(self, key: str, value: Any) -> None:
def search_knowledge_base(self, query: str, k: int = 3) -> List[Document]:
"""
Add or update an entry in the knowledge base.
Searches the knowledge base for the most relevant documents.
Parameters
----------
key : str
The key under which the value will be stored.
value : Any
The value to store. Must be JSON‑serialisable.
Args:
query (str): The search query.
k (int): Number of top documents to return.
Returns:
List[Document]: List of top matching documents.
"""
data = self._load()
data[key] = value
self._save(data)
if not self.documents:
return []
def query_entry(self, key: str) -> Any:
"""
Retrieve the value for a given key.
query_embedding = self.embedding_model.embed_query(query)
query_embedding = np.array(query_embedding).reshape(1, -1)
Parameters
----------
key : str
The key to look up.
similarities = np.dot(self.embeddings, query_embedding.T).flatten()
top_indices = similarities.argsort()[-k:][::-1]
return [self.documents[i] for i in top_indices]
Returns
-------
Any
The stored value.
# Global knowledge base instance
kb = KnowledgeBase()
Raises
------
KeyError
If the key does not exist.
"""
data = self._load()
if key not in data:
raise KeyError(f"Key '{key}' not found in knowledge base.")
return data[key]
def add_to_knowledge_base(content: str) -> str:
"""
Tool wrapper for adding content to the knowledge base.
def delete_entry(self, key: str) -> None:
"""
Delete an entry from the knowledge base.
Args:
content (str): Text to add.
Parameters
----------
key : str
The key to delete.
Returns:
str: Confirmation message.
"""
return kb.add_to_knowledge_base(content)
Raises
------
KeyError
If the key does not exist.
"""
data = self._load()
if key not in data:
raise KeyError(f"Key '{key}' not found in knowledge base.")
del data[key]
self._save(data)
def search_knowledge_base(query: str) -> List[Document]:
"""
Tool wrapper for searching the knowledge base.
def list_entries(self) -> Dict[str, Any]:
"""
Return a copy of all entries in the knowledge base.
Args:
query (str): Search query.
Returns
-------
Dict[str, Any]
All key/value pairs.
"""
return self._load()
Returns:
List[Document]: Matching documents.
"""
return kb.search_knowledge_base(query)