From 42fd00d924228ba37c3511cd9cb616c5baa764d3 Mon Sep 17 00:00:00 2001 From: kuzakhmetovartur Date: Tue, 30 Jun 2026 15:36:47 +0300 Subject: [PATCH] =?UTF-8?q?feat:=20solution=20for=20'=D0=90=D0=B3=D0=B5?= =?UTF-8?q?=D0=BD=D1=82=20=D1=81=20RAG-=D0=BF=D0=B0=D0=BC=D1=8F=D1=82?= =?UTF-8?q?=D1=8C=D1=8E'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 122 ++++++++++++++---------------- src/agent.py | 137 +++++++++++++++------------------- src/embeddings.py | 18 +++++ src/knowledge_base.py | 170 ++++++++++++++++-------------------------- 4 files changed, 198 insertions(+), 249 deletions(-) create mode 100644 src/embeddings.py diff --git a/README.md b/README.md index 6879689..eacd3b6 100644 --- a/README.md +++ b/README.md @@ -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 \ No newline at end of file diff --git a/src/agent.py b/src/agent.py index 3068779..09d7bf7 100644 --- a/src/agent.py +++ b/src/agent.py @@ -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')}") \ No newline at end of file + main() \ No newline at end of file diff --git a/src/embeddings.py b/src/embeddings.py new file mode 100644 index 0000000..33bd57c --- /dev/null +++ b/src/embeddings.py @@ -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) \ No newline at end of file diff --git a/src/knowledge_base.py b/src/knowledge_base.py index 204129c..7e37066 100644 --- a/src/knowledge_base.py +++ b/src/knowledge_base.py @@ -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() \ No newline at end of file + Returns: + List[Document]: Matching documents. + """ + return kb.search_knowledge_base(query) \ No newline at end of file