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

src/knowledge_base.py

@@ -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()