feat: solution for '8. Самописный поисковый агент на основе deep agents from scratch'

src/main.py

@@ -1,60 +1,116 @@
 """
-Command‑line interface for the SearchAgent.
+Main entry point for the deep agent search application.
 
-Usage:
-    python -m src.main "search query here"
-
-The script will print the top 5 results with their relevance scores.
+Demonstrates:
+- Creating a virtual file system.
+- Adding a virtual file with sample data.
+- Performing a search query using a simple deep agent.
 """
 
-import argparse
+from __future__ import annotations
+
 import sys
+from typing import List
 
-from src.agent import SearchAgent
+import numpy as np
+import torch
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+from virtual_file_system import VirtualFileSystem, VirtualFile
 
 
-def main() -> None:
-    parser = argparse.ArgumentParser(description="Deep Agent Search CLI")
-    parser.add_argument(
-        "query",
-        type=str,
-        help="Search query string",
-    )
-    parser.add_argument(
-        "--top",
-        type=int,
-        default=5,
-        help="Number of top results to display (default: 5)",
-    )
-    args = parser.parse_args()
+class SearchAgent:
+    """
+    Simple search agent that ranks lines from a virtual file based on
+    cosine similarity between TF-IDF vectors of the query and the lines.
+    """
 
-    # Example corpus – in a real project this would be loaded from a file
-    corpus = [
-        "Deep learning models can capture complex patterns in data.",
-        "Search engines index documents to provide relevant results.",
-        "PyTorch is a popular deep learning framework.",
-        "Natural language processing involves understanding text.",
-        "Machine learning can be supervised or unsupervised.",
-        "The quick brown fox jumps over the lazy dog.",
-        "Artificial intelligence is transforming many industries.",
-        "Data science combines statistics and programming.",
-        "Neural networks consist of layers of interconnected nodes.",
-        "Optimization algorithms adjust model parameters during training.",
-    ]
+    def __init__(self, vfs: VirtualFileSystem):
+        self.vfs = vfs
 
-    agent = SearchAgent(corpus)
+    def search(self, file_name: str, query: str, top_k: int = 5) -> List[str]:
+        """
+        Search for the most relevant lines in the specified virtual file.
 
-    results = agent.search(args.query, top_k=args.top)
+        Parameters
+        ----------
+        file_name : str
+            Name of the virtual file to search.
+        query : str
+            Search query string.
+        top_k : int, optional
+            Number of top results to return.
 
-    if not results:
-        print("No results found.")
-        sys.exit(0)
+        Returns
+        -------
+        List[str]
+            List of the most relevant lines.
+        """
+        vf = self.vfs.get_file(file_name)
+        data = vf.read().decode("utf-8")
+        lines = [line.strip() for line in data.splitlines() if line.strip()]
+        if not lines:
+            return []
 
-    print(f"Top {len(results)} results for query: '{args.query}'\n")
-    for i, res in enumerate(results, start=1):
-        print(f"{i}. [Doc {res.doc_id}] Score: {res.score:.4f}")
-        print(f"   {res.text}\n")
+        # Vectorize lines and query
+        vectorizer = TfidfVectorizer()
+        doc_vectors = vectorizer.fit_transform(lines).toarray()
+        query_vec = vectorizer.transform([query]).toarray()
+
+        # Convert to torch tensors for similarity calculation
+        doc_tensors = torch.tensor(doc_vectors, dtype=torch.float32)
+        query_tensor = torch.tensor(query_vec, dtype=torch.float32)
+
+        # Normalize vectors
+        doc_norm = doc_tensors / doc_tensors.norm(dim=1, keepdim=True)
+        query_norm = query_tensor / query_tensor.norm()
+
+        # Cosine similarity
+        similarities = torch.matmul(doc_norm, query_norm.t()).squeeze()
+
+        # Get top_k indices
+        top_indices = similarities.topk(top_k).indices.tolist()
+        return [lines[i] for i in top_indices]
+
+
+def main(argv: List[str]) -> None:
+    """
+    Example usage of the virtual file system and search agent.
+
+    Creates a virtual file with sample text and performs a search query.
+    """
+    vfs = VirtualFileSystem()
+
+    # Sample data: a small collection of sentences
+    sample_text = """\
+Deep learning has revolutionized many fields.
+Neural networks can approximate complex functions.
+PyTorch provides dynamic computation graphs.
+Scikit-learn offers a wide range of machine learning tools.
+Numpy is essential for numerical operations.
+"""
+
+    # Create a virtual file
+    vf = vfs.create_file("sample.txt", sample_text.encode("utf-8"))
+
+    # Instantiate the search agent
+    agent = SearchAgent(vfs)
+
+    # Perform a search query
+    query = "neural networks"
+    results = agent.search("sample.txt", query, top_k=3)
+
+    print(f"Search results for query: '{query}'")
+    for idx, line in enumerate(results, 1):
+        print(f"{idx}. {line}")
+
+    # Demonstrate unload
+    vf.unload()
+    try:
+        vf.read()
+    except RuntimeError as e:
+        print(f"After unload: {e}")
 
 
 if __name__ == "__main__":
-    main()
+    main(sys.argv[1:])