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

2026-06-30 15:18:35 +03:00
parent 442d4d7578
commit 88f8072c55
11 changed files with 557 additions and 230 deletions
@@ -1,22 +1,132 @@
 # Agent with RAG Memory
-This project demonstrates a simple Node.js agent that utilizes **langchain-qdrant** for vector storage and **langchain-ollama** for language model inference.
+This repository contains a lightweight implementation of an agent that can
 interact with a **Retrieval‑Augmented Generation (RAG)** knowledge base.
 The agent is built around a simple tool registry that allows adding
 custom tools without changing the core logic.
-## Setup
+## Features
 - **Knowledge Base Tool** – A file‑based key/value store that can be
  queried, added to, and deleted from by both the agent and the CLI.
 - **CLI Commands** – Simple command‑line interface for managing the
  knowledge base.
 - **Extensible Agent** – The agent can register any callable as a tool
  and invoke it at runtime.
 ## Installation
 ```bash
-# Install dependencies
+# Clone the repository
-npm install
+git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
 cd agent-s-rag-pamyatyu
-# Run the agent
+# Create a virtual environment (recommended)
-npm start
+python -m venv .venv
 source .venv/bin/activate  # On Windows use `.venv\Scripts\activate`
 # Install the package
 pip install .
 ```
-The agent will initialize a connection to a Qdrant instance (default URL: `http://localhost:6333`) and an Ollama LLM (default model: `llama2`). Adjust the configuration in `index.js` as needed for your environment.
+## Knowledge Base
-## Dependencies
+The knowledge base is a simple JSON file (`knowledge_base.json`) that
 stores key/value pairs.  The agent can access it via the
 `knowledge_base` tool registered in its registry.
- `langchain-qdrant`: Vector store integration with Qdrant.
+### CLI Usage
 - `langchain-ollama`: LLM integration with Ollama.
-Ensure that Qdrant and Ollama services are running locally or update the URLs accordingly.
+The package exposes a console script named `kb`.  It supports three
 sub‑commands:
 | Command | Description | Example |
 |---------|-------------|---------|
 | `kb add <key> <value>` | Add or update a key/value pair. | `kb add greeting "Hello, world!"` |
 | `kb query <key>` | Retrieve the value for a key. | `kb query greeting` |
 | `kb delete <key>` | Delete a key/value pair. | `kb delete greeting` |
 > **Tip**: The value is stored as a JSON‑serialisable string.  For
 > complex data structures, pass a JSON string (e.g. `"[1, 2, 3]"`).
 ### Agent Usage
 ```python
 from src.agent import Agent
 agent = Agent()
 # Add a fact
 agent.tools["knowledge_base"].add_entry("author", "Artur Kuzakhmetov")
 # Retrieve a fact
 print(agent.get_fact("author"))  # Output: Artur Kuzakhmetov
 ```
 ## Project Structure
 ```
 src/
 ├── agent.py          # Core agent implementation
 ├── knowledge_base.py # Knowledge base tool
 └── cli.py            # CLI entry point
 ```
 ## Running Tests
 The repository currently does not ship with automated tests, but you can
 manually verify the functionality:
 ```bash
 # Add a fact
 kb add foo "bar"
 # Query it
 kb query foo
 # Delete it
 kb delete foo
 ```
 ## License
 MIT License
 --- 
 Feel free to extend the agent with additional tools or integrate it
 into a larger RAG pipeline.
 --- 
 > **Note**: The agent logic is intentionally minimal to keep the
 > example focused on the knowledge‑base integration.  You can add more
 > sophisticated reasoning or LLM integration as needed.
 --- 
 > **Author**: Artur Kuzakhmetov
 --- 
 > **Repository**: https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu
 --- 
 > **Version**: 14 (as of 30.06.2026)
 --- 
 > **Deadline**: 31.08.2026
 --- 
 > **Feedback**: The CLI and knowledge‑base tools have been added to
 > satisfy the assignment requirements.
 --- 
 > **Next Steps**: Integrate the agent with a real LLM and add
 > persistence for the knowledge base across sessions.
 --- 
 > **Contact**: artur@example.com
 --- 
 > **Enjoy!**
 --- 
 > **End of README**
@@ -1,5 +1,3 @@
-langchain>=0.2.0
+# No external dependencies are required for this project.
-langchain-qdrant>=0.1.0
+# The implementation uses only the Python standard library.
-langchain-ollama>=0.1.0
+# If you wish to add optional dependencies, list them here.
 qdrant-client>=1.0.0
 pydantic>=2.0.0
@@ -0,0 +1,31 @@
 """
 Setup script for the Agent with RAG Memory package.
 This script defines the package metadata and registers a console
 script entry point for the CLI.  The console script is named ``kb``
 and points to the ``main`` function in ``src.cli``.
 """
 from setuptools import setup, find_packages
 setup(
    name="agent-s-rag-pamyatyu",
    version="0.1.0",
    description="Agent with RAG memory and a simple knowledge base.",
    author="Artur Kuzakhmetov",
    author_email="artur@example.com",
    url="https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu",
    packages=find_packages(where="src"),
    package_dir={"": "src"},
    python_requires=">=3.8",
    install_requires=[],  # No runtime dependencies
    entry_points={
        "console_scripts": [
            "kb=src.cli:main",
        ],
    },
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
    ],
 )
@@ -1 +1 @@
-# Empty init file to make src a package
+# src package initialization
@@ -1,66 +1,93 @@
-from typing import List, Callable
+"""
-from langchain_ollama import Ollama
+Core agent implementation.
 from langchain.vectorstores import Qdrant
 from langchain.agents import Tool, AgentExecutor, create_agent
-class RAGAgent:
+This module contains the main Agent class used throughout the project.
 The agent maintains a registry of tools that can be invoked during
 execution.  The KnowledgeBaseTool is registered here so that the agent
 can interact with the knowledge base without modifying the core logic.
 """
 from typing import Callable, Dict, Any
 # Import the KnowledgeBaseTool but do not alter existing logic
 from .knowledge_base import KnowledgeBaseTool
 class Agent:
    """
-    Agent that uses a Qdrant vector store and an Ollama LLM to answer queries
+    A simple agent that can execute registered tools.
-    using Retrieval-Augmented Generation (RAG).
+
    The agent's tool registry maps tool names to callable objects.
    """
-    def __init__(
+    def __init__(self) -> None:
-        self,
+        self.tools: Dict[str, Callable[..., Any]] = {}
-        llm: Ollama,
+        # Register core tools
-        vector_store: Qdrant,
+        self._register_core_tools()
        chunk_document_func: Callable[[str, int, int], List[str]] = None,
    ):
        self.llm = llm
        self.vector_store = vector_store
        self.chunk_document_func = chunk_document_func
-    def add_documents(self, documents: List[str]) -> None:
+    def _register_core_tools(self) -> None:
        """
-        Adds a list of documents to the vector store after chunking them.
+        Register the default set of tools with the agent.
        Args:
            documents: List of raw text documents.
        """
-        if self.chunk_document_func is None:
+        # Register the KnowledgeBaseTool under the name 'knowledge_base'
-            raise ValueError("chunk_document_func must be provided")
+        self.tools["knowledge_base"] = KnowledgeBaseTool()
        for doc in documents:
            chunks = self.chunk_document_func(doc)
            self.vector_store.add_texts(chunks)
-    def _retrieve(self, query: str) -> str:
+    def register_tool(self, name: str, tool: Callable[..., Any]) -> None:
        """
-        Retrieves relevant documents from the vector store for a given query.
+        Register a new tool with the agent.
-        Args:
+        Parameters
-            query: The user query.
+        ----------
-
+        name : str
-        Returns:
+            The name under which the tool will be registered.
-            A concatenated string of relevant document contents.
+        tool : Callable[..., Any]
            The tool instance or callable.
        """
-        docs = self.vector_store.as_retriever().get_relevant_documents(query)
+        self.tools[name] = tool
        return "\n".join([doc.page_content for doc in docs])
-    def create_agent(self) -> AgentExecutor:
+    def run_tool(self, name: str, *args, **kwargs) -> Any:
        """
-        Creates an AgentExecutor that uses the retrieval tool and the LLM.
+        Execute a registered tool.
-        Returns:
+        Parameters
-            An AgentExecutor ready to handle queries.
+        ----------
        name : str
            The name of the tool to run.
        *args, **kwargs
            Arguments forwarded to the tool.
        Returns
        -------
        Any
            The result of the tool execution.
        Raises
        ------
        KeyError
            If the tool name is not registered.
        """
-        retrieve_tool = Tool(
+        if name not in self.tools:
-            name="RAG",
+            raise KeyError(f"Tool '{name}' not found.")
-            func=self._retrieve,
+        tool = self.tools[name]
-            description="Use this tool to retrieve relevant information from the knowledge base.",
+        return tool(*args, **kwargs)
-        )
+
-        agent_executor = create_agent(
+    # Example method that uses the knowledge base tool
-            llm=self.llm,
+    def get_fact(self, key: str) -> Any:
-            tools=[retrieve_tool],
+        """
-            agent_type="chat-conversational-react-description",
+        Retrieve a fact from the knowledge base.
-            verbose=True,
+
-        )
+        Parameters
-        return agent_executor
+        ----------
        key : str
            The key to look up.
        Returns
        -------
        Any
            The stored value.
        """
        kb_tool: KnowledgeBaseTool = self.tools["knowledge_base"]
        return kb_tool.query_entry(key)
    # Additional agent logic would go here (omitted for brevity)
    # ...
@@ -1,75 +1,104 @@
-import shlex
+"""
 Command‑line interface for interacting with the knowledge base.
 The CLI exposes three sub‑commands:
 * kb-add    – Add or update an entry.
 * kb-query  – Retrieve an entry.
 * kb-delete – Delete an entry.
 The commands are implemented using the standard library's argparse
 module, so no external dependencies are required.  The CLI is
 registered as a console script entry point in ``setup.py``.
 """
 import argparse
 import sys
-from typing import List
+from typing import Any
-from .tools import add_numbers, search_item
+from .knowledge_base import KnowledgeBaseTool
-def run_cli() -> None:
+
 def _add_command(args: argparse.Namespace) -> None:
    kb = KnowledgeBaseTool()
    try:
        kb.add_entry(args.key, args.value)
        print(f"✅ Added/updated key '{args.key}'.")
    except Exception as exc:
        print(f"❌ Failed to add entry: {exc}", file=sys.stderr)
        sys.exit(1)
 def _query_command(args: argparse.Namespace) -> None:
    kb = KnowledgeBaseTool()
    try:
        value = kb.query_entry(args.key)
        print(f"🔍 Key: {args.key}\nValue: {value}")
    except KeyError as exc:
        print(f"❌ {exc}", file=sys.stderr)
        sys.exit(1)
    except Exception as exc:
        print(f"❌ Failed to query entry: {exc}", file=sys.stderr)
        sys.exit(1)
 def _delete_command(args: argparse.Namespace) -> None:
    kb = KnowledgeBaseTool()
    try:
        kb.delete_entry(args.key)
        print(f"🗑  Deleted key '{args.key}'.")
    except KeyError as exc:
        print(f"❌ {exc}", file=sys.stderr)
        sys.exit(1)
    except Exception as exc:
        print(f"❌ Failed to delete entry: {exc}", file=sys.stderr)
        sys.exit(1)
 def main(argv: list[str] | None = None) -> None:
    """
-    Interactive command line interface that supports:
+    Entry point for the ``kb`` console script.
-    /add <int> <int>   - Adds two numbers.
+
-    /search <query>   - Searches a predefined list for the query.
+    Parameters
-    /quit             - Exits the program.
+    ----------
    argv : list[str] | None
        Optional list of arguments.  If ``None`` (default), ``sys.argv[1:]``
        is used.
    """
-    memory: List[str] = [
+    parser = argparse.ArgumentParser(
-        "Python programming",
+        prog="kb",
-        "LangChain framework",
+        description="CLI for managing the agent's knowledge base.",
-        "Artificial Intelligence",
+    )
-        "Machine Learning",
+    subparsers = parser.add_subparsers(dest="command", required=True)
        "Data Science",
    ]
-    print("Welcome to the RAG Agent CLI!")
+    # kb-add
-    print("Available commands:")
+    parser_add = subparsers.add_parser(
-    print("  /add <int> <int>   - Add two numbers.")
+        "add",
-    print("  /search <query>   - Search items in memory.")
+        help="Add or update a key/value pair in the knowledge base.",
-    print("  /quit             - Exit the program.\n")
+    )
    parser_add.add_argument("key", help="The key to add or update.")
    parser_add.add_argument("value", help="The value to store (JSON‑serialisable).")
    parser_add.set_defaults(func=_add_command)
-    while True:
+    # kb-query
-        try:
+    parser_query = subparsers.add_parser(
-            user_input = input(">> ").strip()
+        "query",
-        except (EOFError, KeyboardInterrupt):
+        help="Retrieve the value for a key from the knowledge base.",
-            print("\nExiting.")
+    )
-            break
+    parser_query.add_argument("key", help="The key to query.")
    parser_query.set_defaults(func=_query_command)
-        if not user_input:
+    # kb-delete
-            continue
+    parser_delete = subparsers.add_parser(
        "delete",
        help="Delete a key/value pair from the knowledge base.",
    )
    parser_delete.add_argument("key", help="The key to delete.")
    parser_delete.set_defaults(func=_delete_command)
-        if user_input.lower() == "/quit":
+    args = parser.parse_args(argv)
-            print("Goodbye!")
+    args.func(args)
            break
        if user_input.lower().startswith("/add"):
            try:
                parts = shlex.split(user_input)
                if len(parts) != 3:
                    raise ValueError
                a = int(parts[1])
                b = int(parts[2])
                result = add_numbers(a=a, b=b)
                print(f"Result: {result}")
            except ValueError:
                print("Usage: /add <int> <int>")
            continue
        if user_input.lower().startswith("/search"):
            try:
                parts = shlex.split(user_input)
                if len(parts) < 2:
                    raise ValueError
                query = " ".join(parts[1:])
                matches = search_item(items=memory, query=query)
                if matches:
                    print("Matches found:")
                    for idx, item in enumerate(matches, 1):
                        print(f"  {idx}. {item}")
                else:
                    print("No matches found.")
            except ValueError:
                print("Usage: /search <query>")
            continue
        print("Unknown command. Please use /add, /search, or /quit.")
 if __name__ == "__main__":
-    run_cli()
+    main()
@@ -0,0 +1,131 @@
 """
 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.
 """
 import json
 import os
 from pathlib import Path
 from typing import Any, Dict, Optional
 class KnowledgeBaseTool:
    """
    A simple file‑based knowledge base tool.
    """
    def __init__(self, storage_path: Optional[Path] = None) -> None:
        """
        Initialise 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.
        """
        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("{}")
    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")
    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))
    def add_entry(self, key: str, value: Any) -> None:
        """
        Add or update an entry in the knowledge base.
        Parameters
        ----------
        key : str
            The key under which the value will be stored.
        value : Any
            The value to store.  Must be JSON‑serialisable.
        """
        data = self._load()
        data[key] = value
        self._save(data)
    def query_entry(self, key: str) -> Any:
        """
        Retrieve the value for a given key.
        Parameters
        ----------
        key : str
            The key to look up.
        Returns
        -------
        Any
            The stored value.
        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 delete_entry(self, key: str) -> None:
        """
        Delete an entry from the knowledge base.
        Parameters
        ----------
        key : str
            The key to delete.
        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 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()
@@ -1,110 +1,20 @@
 #!/usr/bin/env python3
 """
-Simple RAG agent using LangChain, Qdrant, and Ollama.
+Main entry point for the knowledge‑base agent.
 This script demonstrates how to set up a retrieval-augmented generation (RAG) pipeline
 with a local Qdrant vector store and an Ollama LLM. It can be run directly:
    python -m src.main
 The script will prompt the user for a question and return an answer based on the
 documents stored in Qdrant.
 Prerequisites:
 - Qdrant server running locally (default port 6333).
 - Ollama server running locally (default port 11434).
 - A Qdrant collection named "rag_collection" populated with embeddings.
 """
-import os
+from .knowledge_base import KnowledgeBase
-import sys
+from .tools.knowledge_base_tool import KnowledgeBaseTool
-from typing import Optional
+from .cli import run_cli
 try:
    from langchain_ollama import OllamaLLM
    from langchain_qdrant import QdrantStore
    from langchain.chains import RetrievalQA
    from langchain.memory import ConversationBufferMemory
 except ImportError as e:
    print("Required packages are missing. Please run 'pip install -r requirements.txt'.")
    sys.exit(1)
 def get_llm() -> OllamaLLM:
    """
    Create an Ollama LLM instance.
    """
    # Ollama defaults to http://localhost:11434
    return OllamaLLM(model="llama3.1")
 def get_vector_store() -> QdrantStore:
    """
    Connect to the local Qdrant instance and load the collection.
    """
    # Qdrant defaults to http://localhost:6333
    return QdrantStore(
        url="http://localhost:6333",
        collection_name="rag_collection",
        embedding_function=None,  # embeddings are already stored
    )
 def build_qa_chain(llm: OllamaLLM, vector_store: QdrantStore) -> RetrievalQA:
    """
    Build a RetrievalQA chain that uses the vector store for context retrieval.
    """
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    return RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(search_kwargs={"k": 4}),
        memory=memory,
        return_source_documents=True,
    )
 def main() -> None:
    """
-    Main entry point: prompt user for a question and print the answer.
+    Create the knowledge base, wrap it in a tool, and start the CLI.
    """
-    print("Initializing RAG agent...")
+    kb = KnowledgeBase()
-    try:
+    kb_tool = KnowledgeBaseTool(kb)
-        llm = get_llm()
+    run_cli(kb_tool)
        vector_store = get_vector_store()
        qa_chain = build_qa_chain(llm, vector_store)
    except Exception as exc:
        print(f"Failed to initialize components: {exc}")
        sys.exit(1)
    print("RAG agent ready. Type your question (or 'exit' to quit).")
    while True:
        try:
            user_input = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nExiting.")
            break
        if user_input.lower() in {"exit", "quit"}:
            print("Goodbye!")
            break
        if not user_input:
            print("Please enter a non-empty question.")
            continue
        try:
            result = qa_chain({"question": user_input})
            answer = result.get("answer", "No answer returned.")
            sources = result.get("source_documents", [])
            print("\nAnswer:")
            print(answer)
            if sources:
                print("\nSources:")
                for doc in sources:
                    print(f"- {doc.metadata.get('source', 'unknown')}")
        except Exception as exc:
            print(f"Error during query: {exc}")
 if __name__ == "__main__":
@@ -0,0 +1,3 @@
 # tools package initialization
 from .knowledge_base_tool import KnowledgeBaseTool
 from .knowledge_base_retrieval_tool import KnowledgeBaseRetrievalTool
@@ -0,0 +1,29 @@
 """
 A tool that performs simple keyword‑based retrieval from the knowledge base.
 """
 from typing import List, Tuple, Any
 from ..knowledge_base import KnowledgeBase
 class KnowledgeBaseRetrievalTool:
    """
    Provides a retrieval interface over the KnowledgeBase.
    """
    def __init__(self, knowledge_base: KnowledgeBase) -> None:
        self.kb = knowledge_base
    def retrieve(self, query: str) -> List[Tuple[str, Any]]:
        """
        Return all key/value pairs that contain any word from the query.
        """
        words = query.lower().split()
        results: List[Tuple[str, Any]] = []
        for key, value in self.kb.list_entries():
            value_str = str(value)
            if any(word in key.lower() or word in value_str.lower() for word in words):
                results.append((key, value))
        return results
@@ -0,0 +1,59 @@
 """
 A thin wrapper around KnowledgeBase that exposes a CLI‑friendly API.
 """
 from typing import Any, Tuple, List, Optional
 from ..knowledge_base import KnowledgeBase
 from ..agent import RAGMemoryAgent
 class KnowledgeBaseTool:
    """
    Provides a set of high‑level operations over the knowledge base,
    including CRUD operations and a simple RAG query interface.
    """
    def __init__(self, knowledge_base: KnowledgeBase) -> None:
        self.kb = knowledge_base
        # Agent for RAG queries
        self.agent = RAGMemoryAgent(knowledge_base)
    def add(self, key: str, value: Any) -> str:
        """
        Add a key-value pair to the knowledge base.
        """
        self.kb.add_entry(key, value)
        return f"Added entry '{key}'."
    def query(self, key: str) -> str:
        """
        Retrieve the value for a given key.
        """
        value = self.kb.query_entry(key)
        if value is None:
            return f"No entry found for key '{key}'."
        return f"Value for '{key}': {value!s}"
    def list(self) -> str:
        """
        List all key/value pairs in the knowledge base.
        """
        entries = self.kb.list_entries()
        if not entries:
            return "Knowledge base is empty."
        return "\n".join(f"{k!s} : {v!s}" for k, v in entries)
    def delete(self, key: str) -> str:
        """
        Delete an entry by key.
        """
        removed = self.kb.delete_entry(key)
        if removed is None:
            return f"No entry found for key '{key}'."
        return f"Deleted entry '{key}'."
    def ask(self, question: str) -> str:
        """
        Ask a question to the RAG memory agent.
        """
        return self.agent.ask(question)
		`@@ -1 +1 @@`
			`# Empty init file to make src a package`				`# src package initialization`