feat: solution for 'Агент с RAG-памятью'

2026-06-30 15:36:47 +03:00
parent 1279deaaa6
commit 42fd00d924
4 changed files with 198 additions and 249 deletions
@@ -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.  
+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 follows the latest LangChain API and is fully configurable via environment variables.
+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`)
+## Prerequisites
- **Ollama** local LLM integration (via `langchain-ollama`)
+
- Retrieval‑augmented generation with conversation memory
+- Python 3.10+
- Simple CLI interface for quick testing
+- 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)
+# Create a virtual environment
-python -m venv .venv
+python -m venv venv
-source .venv/bin/activate  # On Windows: .venv\\Scripts\\activate
+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):
+Set the Ollama model via environment variable (optional):
 ```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
 ```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
+```bash
-response = agent.run("What is LangChain?")
+export OLLAMA_HOST=http://localhost:11434
 print(response["answer"])
 ```
-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/
+Welcome to the RAG Agent. Type 'exit' to quit.
-├── src/
+User:
 │   ├── agent.py
 │   ├── config.py
 │   └── vector_store.py
 ├── requirements.txt
 ├── pyproject.toml
 └── README.md
 ```
 - **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
@@ -1,94 +1,75 @@
 """
-Agent implementation using the latest LangChain API.
+Agent implementation that uses LangChain to interact with the knowledge base.
 The agent uses a RetrievalQA chain backed by the Qdrant vector store
 and the Ollama LLM for local inference.
 """
-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
+class SearchTool(BaseTool):
 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:
    """
-    A simple retrieval-based agent that answers user queries
+    Tool to search the knowledge base.
    using documents stored in Qdrant and an Ollama LLM.
    """
-
+    name = "search_knowledge_base"
-    def __init__(self, collection_name: str = "documents"):
+    description = (
-        # Initialize LLM
+        "Search the knowledge base for relevant information. "
-        self.llm = Ollama(
+        "Input: query string."
            model=OLLAMA_MODEL,
            base_url=f"http://{OLLAMA_HOST}:{OLLAMA_PORT}",
    )
-        # Initialize vector store
+    def _run(self, query: str) -> str:
-        self.vector_store = QdrantVectorStore(collection_name=collection_name)
+        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])
-        # Memory to keep conversation context
+class AddTool(BaseTool):
-        self.memory = ConversationBufferMemory(memory_key="chat_history")
+    """
-
+    Tool to add new knowledge to the knowledge base.
-        # Prompt template
+    """
-        self.prompt = PromptTemplate(
+    name = "add_to_knowledge_base"
-            input_variables=["chat_history", "question"],
+    description = (
-            template=(
+        "Add new knowledge to the knowledge base. "
-                "You are a helpful assistant. Use the following context to answer the question.\n"
+        "Input: content string."
                "Context:\n{chat_history}\n"
                "Question: {question}\n"
                "Answer:"
            ),
    )
-        # RetrievalQA chain
+    def _run(self, content: str) -> str:
-        self.chain = RetrievalQA.from_chain_type(
+        return add_to_knowledge_base(content)
-            llm=self.llm,
+
-            chain_type="stuff",
+# Instantiate tools
-            retriever=self.vector_store.as_retriever(),
+tools = [SearchTool(), AddTool()]
-            memory=self.memory,
+
-            return_source_documents=True,
+# LLM configuration
-            chain_type_kwargs={"prompt": self.prompt},
+llm = OpenAI(temperature=0)
 # Create the agent
 agent = ZeroShotAgent(llm=llm, tools=tools)
 # Executor that runs the agent
 agent_executor = AgentExecutor.from_agent_and_tools(
    agent=agent,
    tools=tools,
    verbose=True
 )
-    def add_documents(self, documents: list[Document]) -> None:
+def main() -> None:
    """
-        Add documents to the underlying vector store.
+    Simple REPL to interact with the agent.
    """
-        self.vector_store.add_documents(documents)
+    print("Welcome to the RAG Agent. Type 'exit' to quit.")
-
+    while True:
-    def run(self, question: str) -> Any:
+        user_input = input("User: ")
-        """
+        if user_input.lower() in ("exit", "quit"):
-        Run the agent on a user question.
+            print("Goodbye!")
-        Returns the LLM's answer and the source documents.
+            break
-        """
+        try:
-        result = self.chain({"question": question})
+            response = agent_executor.run(user_input)
-        return result
+            print(f"Agent: {response}")
-
+        except Exception as e:
-    def __call__(self, question: str) -> Any:
+            print(f"Error: {e}")
        return self.run(question)
 if __name__ == "__main__":
-    import json
+    main()
    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')}")
@@ -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)
@@ -1,131 +1,89 @@
 """
-Knowledge Base Tool for the Agent.
+Knowledge base implementation using Ollama embeddings.
-
+Provides tools for searching and adding documents.
 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.
 """
-import json
+import numpy as np
-import os
+from typing import List
-from pathlib import Path
+from langchain.docstore.document import Document
-from typing import Any, Dict, Optional
+from embeddings import get_embedding_model
-
+class KnowledgeBase:
 class KnowledgeBaseTool:
    """
-    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
+        Args:
-        ----------
+            content (str): The text content to add.
-        storage_path : Optional[Path]
+
-            Path to the JSON file used for storage.  If not provided,
+        Returns:
-            a file named ``knowledge_base.json`` in the same directory
+            str: Confirmation message.
            as this module is used.
        """
-        if storage_path is None:
+        doc = Document(page_content=content)
-            storage_path = Path(__file__).parent / "knowledge_base.json"
+        embedding = self.embedding_model.embed_query(content)
-        self.storage_path = storage_path
+        embedding = np.array(embedding).reshape(1, -1)
        # Ensure the storage file exists
        if not self.storage_path.exists():
            self.storage_path.write_text("{}")
-    def _load(self) -> Dict[str, Any]:
+        self.documents.append(doc)
-        """Load the knowledge base from disk."""
+        if self.embeddings.size == 0:
-        try:
+            self.embeddings = embedding
-            data = json.loads(self.storage_path.read_text())
+        else:
-            if not isinstance(data, dict):
+            self.embeddings = np.vstack([self.embeddings, embedding])
                raise ValueError("Knowledge base file corrupted: not a dict")
            return data
        except json.JSONDecodeError:
            raise ValueError("Knowledge base file corrupted: invalid JSON")
-    def _save(self, data: Dict[str, Any]) -> None:
+        return f"Document added. Total documents: {len(self.documents)}."
        """Persist the knowledge base to disk."""
        self.storage_path.write_text(json.dumps(data, indent=2, ensure_ascii=False))
-    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
+        Args:
-        ----------
+            query (str): The search query.
-        key : str
+            k (int): Number of top documents to return.
-            The key under which the value will be stored.
+
-        value : Any
+        Returns:
-            The value to store.  Must be JSON‑serialisable.
+            List[Document]: List of top matching documents.
        """
-        data = self._load()
+        if not self.documents:
-        data[key] = value
+            return []
        self._save(data)
-    def query_entry(self, key: str) -> Any:
+        query_embedding = self.embedding_model.embed_query(query)
        query_embedding = np.array(query_embedding).reshape(1, -1)
        similarities = np.dot(self.embeddings, query_embedding.T).flatten()
        top_indices = similarities.argsort()[-k:][::-1]
        return [self.documents[i] for i in top_indices]
 # Global knowledge base instance
 kb = KnowledgeBase()
 def add_to_knowledge_base(content: str) -> str:
    """
-        Retrieve the value for a given key.
+    Tool wrapper for adding content to the knowledge base.
-        Parameters
+    Args:
-        ----------
+        content (str): Text to add.
        key : str
            The key to look up.
-        Returns
+    Returns:
-        -------
+        str: Confirmation message.
        Any
            The stored value.
        Raises
        ------
        KeyError
            If the key does not exist.
    """
-        data = self._load()
+    return kb.add_to_knowledge_base(content)
        if key not in data:
            raise KeyError(f"Key '{key}' not found in knowledge base.")
        return data[key]
-    def delete_entry(self, key: str) -> None:
+def search_knowledge_base(query: str) -> List[Document]:
    """
-        Delete an entry from the knowledge base.
+    Tool wrapper for searching the knowledge base.
-        Parameters
+    Args:
-        ----------
+        query (str): Search query.
        key : str
            The key to delete.
-        Raises
+    Returns:
-        ------
+        List[Document]: Matching documents.
        KeyError
            If the key does not exist.
    """
-        data = self._load()
+    return kb.search_knowledge_base(query)
        if key not in data:
            raise KeyError(f"Key '{key}' not found in knowledge base.")
        del data[key]
        self._save(data)
    def list_entries(self) -> Dict[str, Any]:
        """
        Return a copy of all entries in the knowledge base.
        Returns
        -------
        Dict[str, Any]
            All key/value pairs.
        """
        return self._load()