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

src/index.py

@@ -1,92 +1,110 @@
-#!/usr/bin/env python3
-"""
-Simple search agent implementation.
-
-This module provides a minimal command‑line interface that accepts a search
-query and returns a list of dummy results.  It is intentionally lightweight
-to satisfy the assignment requirements while demonstrating a clear
-structure that can be expanded in the future.
-
-Author: Artur Kuzakhmetov
-"""
-
-import argparse
-import sys
+import os
+import json
+import numpy as np
+import torch
+from transformers import AutoTokenizer, AutoModel
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from typing import List
 
+class SearchRequest(BaseModel):
+    query: str
+    top_k: int = 5
 
-def search(query: str, limit: int = 5) -> List[str]:
+class SearchResult(BaseModel):
+    document: str
+    score: float
+
+class SearchResponse(BaseModel):
+    results: List[SearchResult]
+
+class SearchAgent:
     """
-    Perform a mock search for the given query.
-
-    Parameters
-    ----------
-    query : str
-        The search string.
-    limit : int, optional
-        Maximum number of results to return. Defaults to 5.
-
-    Returns
-    -------
-    List[str]
-        A list of fake search results.
-
-    Notes
-    -----
-    This function does not perform real network requests.  It simply
-    generates deterministic placeholder results so that the module can be
-    tested without external dependencies.
+    A simple deep‑agent search engine that uses a transformer encoder
+    to embed documents and queries, then ranks documents by cosine
+    similarity.
     """
-    if not query:
-        raise ValueError("Query must not be empty")
+    def __init__(self,
+                 data_path: str = "data/documents.txt",
+                 model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name).to(self.device)
+        self.documents = self._load_documents(data_path)
+        self.embeddings = self._embed_documents(self.documents)
 
-    # Generate deterministic dummy results
-    results = [f"{query} result {i+1}" for i in range(limit)]
-    return results
+    def _load_documents(self, path: str) -> List[str]:
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Data file not found: {path}")
+        with open(path, "r", encoding="utf-8") as f:
+            docs = [line.strip() for line in f if line.strip()]
+        return docs
 
+    def _embed_documents(self, docs: List[str]) -> np.ndarray:
+        batch_size = 32
+        embeddings = []
+        for i in range(0, len(docs), batch_size):
+            batch = docs[i:i+batch_size]
+            inputs = self.tokenizer(batch,
+                                    padding=True,
+                                    truncation=True,
+                                    return_tensors="pt").to(self.device)
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                token_embeddings = outputs.last_hidden_state
+                attention_mask = inputs["attention_mask"].unsqueeze(-1)
+                sum_embeddings = torch.sum(token_embeddings * attention_mask, dim=1)
+                sum_mask = torch.clamp(attention_mask.sum(dim=1), min=1e-9)
+                batch_embeddings = sum_embeddings / sum_mask
+                embeddings.append(batch_embeddings.cpu().numpy())
+        return np.vstack(embeddings)
 
-def main(argv: List[str] | None = None) -> int:
+    def _embed_query(self, query: str) -> np.ndarray:
+        inputs = self.tokenizer([query],
+                                padding=True,
+                                truncation=True,
+                                return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            token_embeddings = outputs.last_hidden_state
+            attention_mask = inputs["attention_mask"].unsqueeze(-1)
+            sum_embeddings = torch.sum(token_embeddings * attention_mask, dim=1)
+            sum_mask = torch.clamp(attention_mask.sum(dim=1), min=1e-9)
+            query_embedding = sum_embeddings / sum_mask
+            return query_embedding.cpu().numpy()
+
+    def search(self, query: str, top_k: int = 5) -> List[SearchResult]:
+        query_emb = self._embed_query(query)
+        dot = np.dot(self.embeddings, query_emb.T).squeeze()
+        norms = np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_emb)
+        similarities = dot / norms
+        top_indices = np.argsort(similarities)[::-1][:top_k]
+        results = [SearchResult(document=self.documents[idx],
+                                score=float(similarities[idx]))
+                   for idx in top_indices]
+        return results
+
+app = FastAPI(title="Deep Agent Search API")
+
+# Instantiate the agent once at startup
+agent = SearchAgent()
+
+@app.post("/search", response_model=SearchResponse)
+async def search_endpoint(request: SearchRequest):
     """
-    Entry point for the command‑line interface.
-
-    Parameters
-    ----------
-    argv : List[str] | None
-        List of command‑line arguments.  If None, sys.argv[1:] is used.
-
-    Returns
-    -------
-    int
-        Exit code (0 for success, 1 for error).
+    Search endpoint that accepts a JSON payload:
+    {
+        "query": "your search query",
+        "top_k": 5
+    }
+    Returns the top_k most relevant documents.
     """
-    parser = argparse.ArgumentParser(
-        description="Simple search agent – returns mock results for a query."
-    )
-    parser.add_argument(
-        "query",
-        type=str,
-        help="Search query string",
-    )
-    parser.add_argument(
-        "-n",
-        "--limit",
-        type=int,
-        default=5,
-        help="Number of results to return (default: 5)",
-    )
-    args = parser.parse_args(argv)
-
     try:
-        results = search(args.query, args.limit)
-    except ValueError as exc:
-        print(f"Error: {exc}", file=sys.stderr)
-        return 1
-
-    for idx, result in enumerate(results, start=1):
-        print(f"{idx}. {result}")
-
-    return 0
-
+        results = agent.search(request.query, request.top_k)
+        return SearchResponse(results=results)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    sys.exit(main())
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)