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

src/index.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+"""
+Educational Agent with Retrieval-Augmented Generation (RAG) Memory.
+
+This module implements a lightweight RAG system using only Python standard
+libraries. It provides:
+
+* RAGMemory – stores documents, builds simple bag‑of‑words embeddings,
+  and retrieves the most relevant passages for a query.
+* Agent – integrates the memory with a rule‑based response generator.
+* CLI – allows the user to ask a question and receive an answer that
+  incorporates retrieved context.
+
+The implementation avoids external AI services or heavy dependencies,
+making it suitable for the Deep Agents Virtual File System environment.
+"""
+
+import os
+import sys
+import math
+import argparse
+from collections import defaultdict
+from typing import Dict, List, Tuple
+
+# --------------------------------------------------------------------------- #
+# Utility functions
+# --------------------------------------------------------------------------- #
+
+def tokenize(text: str) -> List[str]:
+    """
+    Very simple tokenizer: lowercases, splits on whitespace, removes
+    punctuation.
+    """
+    import string
+    translator = str.maketrans('', '', string.punctuation)
+    return text.translate(translator).lower().split()
+
+def vectorize(tokens: List[str]) -> Dict[str, int]:
+    """
+    Convert a list of tokens into a bag‑of‑words vector (word -> count).
+    """
+    vec = defaultdict(int)
+    for token in tokens:
+        vec[token] += 1
+    return dict(vec)
+
+def dot_product(v1: Dict[str, int], v2: Dict[str, int]) -> int:
+    """
+    Compute dot product of two sparse vectors represented as dicts.
+    """
+    if len(v1) > len(v2):
+        v1, v2 = v2, v1
+    return sum(v1.get(k, 0) * v2.get(k, 0) for k in v1)
+
+def vector_norm(v: Dict[str, int]) -> float:
+    """
+    Compute Euclidean norm of a sparse vector.
+    """
+    return math.sqrt(sum(count * count for count in v.values()))
+
+def cosine_similarity(v1: Dict[str, int], v2: Dict[str, int]) -> float:
+    """
+    Compute cosine similarity between two sparse vectors.
+    """
+    denom = vector_norm(v1) * vector_norm(v2)
+    if denom == 0:
+        return 0.0
+    return dot_product(v1, v2) / denom
+
+# --------------------------------------------------------------------------- #
+# RAG Memory
+# --------------------------------------------------------------------------- #
+
+class RAGMemory:
+    """
+    Simple RAG memory that stores documents and their embeddings.
+    Retrieval is performed via cosine similarity over bag‑of‑words vectors.
+    """
+
+    def __init__(self):
+        # doc_id -> (text, vector)
+        self._store: Dict[str, Tuple[str, Dict[str, int]]] = {}
+
+    def add_document(self, doc_id: str, text: str) -> None:
+        """
+        Add a document to the memory.
+        """
+        tokens = tokenize(text)
+        vec = vectorize(tokens)
+        self._store[doc_id] = (text, vec)
+
+    def load_from_directory(self, path: str) -> None:
+        """
+        Load all .txt files from a directory as documents.
+        """
+        for filename in os.listdir(path):
+            if filename.lower().endswith('.txt'):
+                doc_id = os.path.splitext(filename)[0]
+                with open(os.path.join(path, filename), 'r', encoding='utf-8') as f:
+                    text = f.read()
+                self.add_document(doc_id, text)
+
+    def retrieve(self, query: str, top_k: int = 3) -> List[Tuple[str, str, float]]:
+        """
+        Retrieve top_k documents most relevant to the query.
+        Returns a list of tuples: (doc_id, text, similarity_score).
+        """
+        query_vec = vectorize(tokenize(query))
+        scores = []
+        for doc_id, (text, vec) in self._store.items():
+            score = cosine_similarity(query_vec, vec)
+            scores.append((doc_id, text, score))
+        # Sort by descending similarity
+        scores.sort(key=lambda x: x[2], reverse=True)
+        return scores[:top_k]
+
+# --------------------------------------------------------------------------- #
+# Agent
+# --------------------------------------------------------------------------- #
+
+class Agent:
+    """
+    Educational agent that uses RAGMemory to augment its responses.
+    """
+
+    def __init__(self, memory: RAGMemory):
+        self.memory = memory
+
+    def answer(self, question: str, top_k: int = 3) -> str:
+        """
+        Generate an answer to the question by retrieving context and
+        concatenating it with a simple rule‑based response.
+        """
+        retrieved = self.memory.retrieve(question, top_k=top_k)
+        context_parts = []
+        for doc_id, text, score in retrieved:
+            snippet = text[:200].replace('\n', ' ')  # short snippet
+            context_parts.append(f"[{doc_id}] {snippet} (score={score:.2f})")
+        context = "\n".join(context_parts) if context_parts else "No relevant context found."
+        answer = (
+            f"Question: {question}\n\n"
+            f"Context:\n{context}\n\n"
+            f"Answer: (This is a placeholder answer generated by a rule‑based system.)"
+        )
+        return answer
+
+# --------------------------------------------------------------------------- #
+# CLI
+# --------------------------------------------------------------------------- #
+
+def main():
+    parser = argparse.ArgumentParser(description="Educational Agent with RAG Memory")
+    parser.add_argument("--data-dir", type=str, required=True,
+                        help="Directory containing .txt documents for the memory.")
+    parser.add_argument("--question", type=str, required=True,
+                        help="The question to ask the agent.")
+    parser.add_argument("--top-k", type=int, default=3,
+                        help="Number of top documents to retrieve.")
+    args = parser.parse_args()
+
+    memory = RAGMemory()
+    memory.load_from_directory(args.data_dir)
+
+    agent = Agent(memory)
+    response = agent.answer(args.question, top_k=args.top_k)
+    print(response)
+
+if __name__ == "__main__":
+    main()