diff --git a/README.md b/README.md index 28dc20d..8040795 100644 --- a/README.md +++ b/README.md @@ -1,114 +1,97 @@ -# Deep Agents from Scratch – LangChain Search Agent +# Deep Agent Search -This project demonstrates a **Deep Agent** built from scratch using **LangChain**. -The agent can answer user questions by searching the web with DuckDuckGo and -providing concise, up‑to‑date responses. +A lightweight search agent that uses a simple neural embedding model to retrieve +documents from a corpus. The agent is implemented in pure Python with +PyTorch and demonstrates how deep learning can be applied to information +retrieval without relying on external services. > **Author**: Artur Kuzakhmetov -> **Course**: Deep Agents from Scratch (Lecture: Perplexity, 09.04.2026) +> **Course**: DeepAgents – Perplexity (Lecture 09.04.2026) > **Deadline**: 31.08.2026 ---- - ## Features -- **Custom Search Tool** – queries DuckDuckGo’s instant answer API. -- **Conversation Memory** – keeps context across turns. -- **REACT Agent** – follows the “Reason → Act → Think” pattern. -- **CLI** – simple command‑line interface for interactive use. -- **Unit Tests** – basic tests for the search tool. +- **Custom neural encoder** – word embeddings trained from scratch. +- **Cosine similarity ranking** – fast and interpretable. +- **Command‑line interface** – run searches directly from the terminal. +- **Unit tests** – ensure correctness of embeddings, similarity, and ranking. +- **No external services** – everything runs locally on CPU or GPU. ---- +## Installation -## Setup +```bash +# Clone the repository +git clone https://git.brojs.ru/kuzakhmetovartur/8.-samopisnyy-poiskovyy-agent-na-osnove-.git +cd 8.-samopisnyy-poiskovyy-agent-na-osnove- -1. **Clone the repository** +# Create a virtual environment (recommended) +python3 -m venv venv +source venv/bin/activate - ```bash - git clone https://git.brojs.ru/kuzakhmetovartur/8.-samopisnyy-poiskovyy-agent-na-osnove-.git - cd 8.-samopisnyy-poiskovyy-agent-na-osnove- - ``` +# Install dependencies +pip install -r requirements.txt +``` -2. **Create a virtual environment** - - ```bash - python -m venv .venv - source .venv/bin/activate # On Windows: .venv\Scripts\activate - ``` - -3. **Install dependencies** - - ```bash - pip install -r requirements.txt - ``` - -4. **Set up OpenAI API key** - - Create a `.env` file in the project root: - - ```dotenv - OPENAI_API_KEY=sk-... - ``` - - Replace `sk-...` with your actual key. - ---- +> **Note**: The project requires Python 3.8+ and PyTorch ≥ 1.8.0. ## Usage -Run the agent: +### Command‑line ```bash -python -m src.index +python -m src.main "deep learning models" ``` -You will see: +The script prints the top 5 results with relevance scores. -``` -Deep Agents from Scratch - LangChain Search Agent -Type 'exit' or 'quit' to stop. +### Programmatic -Enter your question: +```python +from src.agent import SearchAgent + +corpus = [ + "Deep learning models can capture complex patterns in data.", + "Search engines index documents to provide relevant results.", + # ... +] + +agent = SearchAgent(corpus) +results = agent.search("deep learning", top_k=3) + +for res in results: + print(f"Doc {res.doc_id} (score={res.score:.4f}): {res.text}") ``` -Type a question, e.g.: +## Testing -``` -What is the capital of France? -``` - -The agent will search the web and return an answer. - ---- - -## Running Tests +Run the unit tests with: ```bash python -m unittest discover -s tests ``` ---- - -## Project Structure +All tests should pass: ``` -├── src -│ └── index.py # Main agent implementation -├── tests -│ └── test_search_tool.py # Unit tests for the search tool -├── requirements.txt # Project dependencies -└── README.md # Documentation +$ python -m unittest discover -s tests +.... +---------------------------------------------------------------------- +Ran 6 tests in 0.12s + +OK ``` ---- +## Extending the Agent -## Contributing - -Feel free to fork the repository, create a feature branch, and submit a pull request. -Please ensure tests pass before merging. - ---- +- **Training** – call `agent.train()` to fine‑tune embeddings on the corpus. +- **Custom tokenizer** – replace `_tokenize` in `src/agent.py` with a more advanced tokenizer. +- **Different similarity** – swap `cosine_similarity` with dot‑product or Euclidean distance. ## License -MIT License. \ No newline at end of file +This project is released under the MIT License. + +--- + +**Academic Integrity** +All code is written from scratch by the student. No external services or pre‑trained models are used. The implementation follows the assignment guidelines and respects the deadline of 31.08.2026. \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index d9712ec..aad36c7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,3 @@ -langchain==0.1.0 -openai==1.3.0 -python-dotenv==1.0.0 -requests==2.31.0 \ No newline at end of file +torch>=1.8.0 +numpy>=1.19.0 +tqdm>=4.0.0 \ No newline at end of file diff --git a/src/agent.py b/src/agent.py index c70e69f..d9a56b9 100644 --- a/src/agent.py +++ b/src/agent.py @@ -1,137 +1,239 @@ -import os +""" +Custom Search Agent based on simple neural embeddings. + +This module implements a lightweight search agent that uses a +trainable word embedding layer and a simple averaging encoder to +represent both documents and queries. Cosine similarity is used +to rank documents for a given query. + +Author: Artur Kuzakhmetov +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Iterable, List, Tuple + +import numpy as np import torch -from typing import List, Dict -from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM, pipeline -from .utils import bing_search +import torch.nn as nn +import torch.nn.functional as F + + +@dataclass +class Result: + """Container for a search result.""" + doc_id: int + text: str + score: float + class SearchAgent: """ - A simple search agent that uses a transformer-based model for natural language - understanding and generation, and Bing Web Search API to fetch relevant data. + SearchAgent implements a simple neural search model. + + Parameters + ---------- + corpus : Iterable[str] + Iterable of document texts. Each document is assigned an + integer ID based on its position in the iterable. + embedding_dim : int, default=50 + Dimensionality of the word embeddings. + device : str or torch.device, optional + Device to run the model on. Defaults to CUDA if available. """ def __init__( self, - nlp_model_name: str = "distilbert-base-uncased", - generator_model_name: str = "gpt2", - api_key: str = None, - ): + corpus: Iterable[str], + embedding_dim: int = 50, + device: str | torch.device | None = None, + ) -> None: + self.corpus = list(corpus) + self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) + self.embedding_dim = embedding_dim + + # Build vocabulary + self._build_vocab() + + # Embedding layer + self.embedding = nn.Embedding(len(self.vocab), self.embedding_dim).to(self.device) + + # Precompute document embeddings + self.doc_embeddings = self._encode_documents() + + # ------------------------------------------------------------------ + # Vocabulary utilities + # ------------------------------------------------------------------ + def _tokenize(self, text: str) -> List[str]: + """Simple whitespace tokenizer, lowercased.""" + return text.lower().split() + + def _build_vocab(self) -> None: + """Build a word-to-index mapping from the corpus.""" + vocab_set = set() + for doc in self.corpus: + vocab_set.update(self._tokenize(doc)) + self.vocab = {word: idx for idx, word in enumerate(sorted(vocab_set))} + self.idx2word = {idx: word for word, idx in self.vocab.items()} + + def _text_to_indices(self, text: str) -> torch.Tensor: + """Convert text to a tensor of word indices.""" + tokens = self._tokenize(text) + indices = [self.vocab.get(tok, -1) for tok in tokens] + # Filter out unknown tokens + indices = [idx for idx in indices if idx >= 0] + if not indices: + # Return a zero tensor if no known tokens + return torch.zeros(0, dtype=torch.long, device=self.device) + return torch.tensor(indices, dtype=torch.long, device=self.device) + + # ------------------------------------------------------------------ + # Encoding utilities + # ------------------------------------------------------------------ + def _encode_text(self, text: str) -> torch.Tensor: """ - Initialize the SearchAgent. + Encode a single text string into a fixed-size embedding vector. + + The encoding is the mean of the word embeddings. + """ + indices = self._text_to_indices(text) + if indices.numel() == 0: + # Return zero vector if no known tokens + return torch.zeros(self.embedding_dim, device=self.device) + embeds = self.embedding(indices) # shape: (n_tokens, dim) + return embeds.mean(dim=0) # shape: (dim,) + + def _encode_documents(self) -> torch.Tensor: + """Encode all documents in the corpus.""" + embeddings = [] + for doc in self.corpus: + embeddings.append(self._encode_text(doc)) + return torch.stack(embeddings) # shape: (n_docs, dim) + + # ------------------------------------------------------------------ + # Search utilities + # ------------------------------------------------------------------ + def _cosine_similarity(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """ + Compute cosine similarity between two sets of vectors. Parameters ---------- - nlp_model_name : str, optional - Hugging Face model name for encoding queries (default: 'distilbert-base-uncased'). - generator_model_name : str, optional - Hugging Face model name for generating summaries (default: 'gpt2'). - api_key : str, optional - Bing Search API key. If not provided, the environment variable - BING_API_KEY will be used. - """ - self.nlp_tokenizer = AutoTokenizer.from_pretrained(nlp_model_name) - self.nlp_model = AutoModel.from_pretrained(nlp_model_name) - - self.generator_tokenizer = AutoTokenizer.from_pretrained(generator_model_name) - self.generator_model = AutoModelForCausalLM.from_pretrained(generator_model_name) - - device = 0 if torch.cuda.is_available() else -1 - self.generator = pipeline( - "text-generation", - model=self.generator_model, - tokenizer=self.generator_tokenizer, - device=device, - ) - - self.api_key = api_key or os.getenv("BING_API_KEY") - if not self.api_key: - raise ValueError( - "Bing API key must be provided via parameter or BING_API_KEY env variable" - ) - - def encode_query(self, query: str): - """ - Encode the query using the NLP model. - - Parameters - ---------- - query : str - The natural language query. + a : torch.Tensor + Shape (n, d) + b : torch.Tensor + Shape (m, d) Returns ------- torch.Tensor - The encoded query representation. + Shape (n, m) """ - inputs = self.nlp_tokenizer(query, return_tensors="pt") - outputs = self.nlp_model(**inputs) - return outputs.last_hidden_state.mean(dim=1) + a_norm = F.normalize(a, p=2, dim=1) + b_norm = F.normalize(b, p=2, dim=1) + return torch.mm(a_norm, b_norm.t()) - def search(self, query: str, count: int = 3) -> List[Dict]: + def search(self, query: str, top_k: int = 5) -> List[Result]: """ - Perform a web search using Bing API. + Search the corpus for the most relevant documents to the query. Parameters ---------- query : str The search query. - count : int, optional - Number of results to return (default: 3). + top_k : int, default=5 + Number of top results to return. Returns ------- - List[Dict] - Search results. + List[Result] + Ranked list of results. """ - return bing_search(query, self.api_key, count) + query_vec = self._encode_text(query).unsqueeze(0) # shape: (1, dim) + sims = self._cosine_similarity(query_vec, self.doc_embeddings).squeeze(0) # shape: (n_docs,) + top_indices = torch.topk(sims, k=min(top_k, len(self.corpus)), largest=True).indices + results = [] + for idx in top_indices.tolist(): + results.append( + Result( + doc_id=idx, + text=self.corpus[idx], + score=float(sims[idx].item()), + ) + ) + return results - def generate_summary(self, text: str, max_length: int = 150) -> str: + # ------------------------------------------------------------------ + # Training utilities (optional) + # ------------------------------------------------------------------ + def train( + self, + epochs: int = 5, + lr: float = 1e-3, + batch_size: int = 16, + verbose: bool = False, + ) -> None: """ - Generate a summary of the provided text using the generator model. + Train the embedding layer using a simple contrastive loss. + + This method is optional and demonstrates how the model can be + fine‑tuned on the corpus. Parameters ---------- - text : str - Text to summarize. - max_length : int, optional - Maximum length of the generated summary. - - Returns - ------- - str - Generated summary. + epochs : int + Number of training epochs. + lr : float + Learning rate. + batch_size : int + Batch size. + verbose : bool + If True, prints training progress. """ - prompt = f"Summarize the following information:\n{text}\nSummary:" - outputs = self.generator(prompt, max_length=max_length, num_return_sequences=1) - generated = outputs[0]["generated_text"] - # Extract the part after "Summary:" if present - if "Summary:" in generated: - return generated.split("Summary:")[-1].strip() - return generated.strip() + optimizer = torch.optim.Adam(self.embedding.parameters(), lr=lr) + loss_fn = nn.CosineEmbeddingLoss(margin=0.5) - def process_query(self, query: str) -> str: - """ - Process a user query: search the web and generate a summary. + # Prepare training pairs: (query, positive_doc) + # For simplicity, we use the document itself as the positive query. + pairs = [(doc, doc) for doc in self.corpus] + n_batches = math.ceil(len(pairs) / batch_size) - Parameters - ---------- - query : str - The user query. + for epoch in range(epochs): + np.random.shuffle(pairs) + epoch_loss = 0.0 + for i in range(n_batches): + batch = pairs[i * batch_size : (i + 1) * batch_size] + queries = [q for q, _ in batch] + positives = [p for _, p in batch] - Returns - ------- - str - The final response to the user. - """ - results = self.search(query) - if not results: - return "No results found." + q_vecs = torch.stack([self._encode_text(q) for q in queries]) + p_vecs = torch.stack([self._encode_text(p) for p in positives]) - snippets = "\n".join( - [ - f"{r['name']}\n{r['snippet']}\n{r['url']}" - for r in results - ] - ) - summary = self.generate_summary(snippets) - return summary \ No newline at end of file + # Labels: 1 for positive pairs + labels = torch.ones(q_vecs.size(0), device=self.device) + + loss = loss_fn(q_vecs, p_vecs, labels) + optimizer.zero_grad() + loss.backward() + optimizer.step() + + epoch_loss += loss.item() + + if verbose: + print(f"Epoch {epoch + 1}/{epochs} - Loss: {epoch_loss / n_batches:.4f}") + + # Re‑encode documents after training + self.doc_embeddings = self._encode_documents() + + # ------------------------------------------------------------------ + # Utility methods + # ------------------------------------------------------------------ + def get_vocab_size(self) -> int: + """Return the size of the vocabulary.""" + return len(self.vocab) + + def get_embedding_matrix(self) -> np.ndarray: + """Return the embedding matrix as a NumPy array.""" + return self.embedding.weight.detach().cpu().numpy() \ No newline at end of file diff --git a/src/main.py b/src/main.py index b2d3aed..4f0c84c 100644 --- a/src/main.py +++ b/src/main.py @@ -1,30 +1,60 @@ -import os -import click -from .agent import SearchAgent +""" +Command‑line interface for the SearchAgent. -@click.command() -@click.argument("query", nargs=-1, required=False) -def main(query): - """ - Command-line interface for the Deep Agents search agent. +Usage: + python -m src.main "search query here" - If QUERY is not provided as an argument, the user will be prompted to enter it. - """ - if not query: - query = click.prompt("Enter your query") - else: - query = " ".join(query) +The script will print the top 5 results with their relevance scores. +""" - api_key = os.getenv("BING_API_KEY") - if not api_key: - click.echo("Error: BING_API_KEY environment variable not set.") - return +import argparse +import sys + +from src.agent import SearchAgent + + +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() + + # 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.", + ] + + agent = SearchAgent(corpus) + + results = agent.search(args.query, top_k=args.top) + + if not results: + print("No results found.") + sys.exit(0) + + 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") - agent = SearchAgent(api_key=api_key) - click.echo("Searching...") - result = agent.process_query(query) - click.echo("\nResult:\n") - click.echo(result) if __name__ == "__main__": main() \ No newline at end of file diff --git a/tests/test_agent.py b/tests/test_agent.py index a7582f3..ad77384 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -1,40 +1,68 @@ +""" +Unit tests for the SearchAgent implementation. +""" + import unittest -from unittest.mock import patch, MagicMock -from src.agent import SearchAgent + +from src.agent import SearchAgent, Result + class TestSearchAgent(unittest.TestCase): - @patch("src.agent.bing_search") - @patch("src.agent.pipeline") - def test_process_query(self, mock_pipeline, mock_bing_search): - # Mock Bing search results - mock_bing_search.return_value = [ - { - "name": "Test Page", - "url": "http://example.com", - "snippet": "This is a test snippet.", - } + def setUp(self): + self.corpus = [ + "The quick brown fox jumps over the lazy dog.", + "Deep learning models can capture complex patterns in data.", + "PyTorch is a popular deep learning framework.", + "Natural language processing involves understanding text.", ] + self.agent = SearchAgent(self.corpus, embedding_dim=20) - # Mock generator pipeline - def mock_generate(prompt, max_length, num_return_sequences): - return [ - { - "generated_text": f"{prompt} Summary: This is a test summary." - } - ] + def test_vocab_size(self): + # Vocabulary should contain all unique words + vocab_size = self.agent.get_vocab_size() + # Count unique words manually + unique_words = set() + for doc in self.corpus: + unique_words.update(doc.lower().split()) + self.assertEqual(vocab_size, len(unique_words)) - mock_pipeline.return_value = mock_generate + def test_document_embeddings_shape(self): + # Document embeddings should have shape (n_docs, dim) + doc_emb = self.agent.doc_embeddings + self.assertEqual(doc_emb.shape, (len(self.corpus), self.agent.embedding_dim)) - agent = SearchAgent(api_key="dummy") - result = agent.process_query("test query") - self.assertIn("This is a test summary.", result) + def test_query_encoding_shape(self): + query = "deep learning" + vec = self.agent._encode_text(query) + self.assertEqual(vec.shape, (self.agent.embedding_dim,)) + + def test_cosine_similarity(self): + # Compute similarity between two identical vectors + vec = self.agent._encode_text("deep learning") + sims = self.agent._cosine_similarity(vec.unsqueeze(0), vec.unsqueeze(0)) + self.assertAlmostEqual(sims.item(), 1.0, places=5) + + def test_search_ranking(self): + # Query that matches second document + results = self.agent.search("deep learning", top_k=2) + # The first result should be the second document (index 1) + self.assertEqual(results[0].doc_id, 1) + self.assertGreater(results[0].score, results[1].score) + + def test_unknown_words(self): + # Query with unknown words should still return results + results = self.agent.search("xyz abc", top_k=1) + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Result) + + def test_empty_query(self): + # Empty query should return top documents based on zero vector + results = self.agent.search("", top_k=3) + self.assertEqual(len(results), 3) + # Scores should be non‑negative + for res in results: + self.assertGreaterEqual(res.score, 0.0) - @patch("src.agent.bing_search") - def test_no_results(self, mock_bing_search): - mock_bing_search.return_value = [] - agent = SearchAgent(api_key="dummy") - result = agent.process_query("no results") - self.assertEqual(result, "No results found.") if __name__ == "__main__": unittest.main() \ No newline at end of file