kuzakhmetovartur/8.-samopisnyy-poiskovyy-agent-na-osnove-
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: solution for '8. Самописный поисковый агент на основе deep agents from scratch'
CI / build (3.1) (push) Has been cancelled
Details
CI / build (3.11) (push) Has been cancelled
Details
CI / build (3.8) (push) Has been cancelled
Details
CI / build (3.9) (push) Has been cancelled
Details

Browse Source
This commit is contained in:
Артур Кузахметов
2026-06-30 16:22:52 +03:00
parent df9e4f49d2
commit 1039c7065c
14 changed files with 524 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+33
View File
@@ -0,0 +1,33 @@
name: CI
on:
push:
branches: [ main ]
pull_request:
branches: [ main ]
jobs:
build:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: [3.8, 3.9, 3.10, 3.11]
steps:
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install -e .
pip install pytest coverage
- name: Run tests
run: |
pytest --maxfail=1 --disable-warnings -q
- name: Upload coverage
uses: actions/upload-artifact@v4
with:
name: coverage-${{ matrix.python-version }}
path: .coverage
.gitignore
+62 -4
View File
@@ -1,5 +1,63 @@
node_modules/
.env
dist/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
*.log
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to keep the file
# from being overwritten.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
# Jupyter Notebook
.ipynb_checkpoints
# VS Code
.vscode/
# Virtual environment
.venv/
env/
ENV/
env.bak/
venv/
# End of file
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 Your Name
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the “Software”), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
README.md
+20 -74
View File
@@ -1,97 +1,43 @@
# Deep Agent Search
# DeepAgent
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**: DeepAgents – Perplexity (Lecture 09.04.2026)
> **Deadline**: 31.08.2026
## Features
- **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.
DeepAgent is a minimal example of a deep learning based search agent.
It demonstrates how to combine a neural network with a simple search algorithm
(Monte‑Carlo Tree Search style) without relying on external search libraries.
## Installation
```bash
# Clone the repository
git clone https://git.brojs.ru/kuzakhmetovartur/8.-samopisnyy-poiskovyy-agent-na-osnove-<repo>.git
cd 8.-samopisnyy-poiskovyy-agent-na-osnove-<repo>
# Create a virtual environment (recommended)
python3 -m venv venv
source venv/bin/activate
python -m venv .venv
source .venv/bin/activate # On Windows use `.venv\\Scripts\\activate`
# Install dependencies
pip install -r requirements.txt
# Install the package
pip install .
```
> **Note**: The project requires Python 3.8+ and PyTorch ≥ 1.8.0.
## Usage
### Command‑line
```bash
python -m src.main "deep learning models"
```
The script prints the top 5 results with relevance scores.
### Programmatic
```python
from src.agent import SearchAgent
from src.search_agent import SearchAgent, PolicyValueNet
corpus = [
"Deep learning models can capture complex patterns in data.",
"Search engines index documents to provide relevant results.",
# ...
]
# Create a policy‑value network
net = PolicyValueNet(input_dim=1, action_space=2)
agent = SearchAgent(corpus)
results = agent.search("deep learning", top_k=3)
# Create the agent
agent = SearchAgent(policy_value_net=net, max_depth=3)
for res in results:
print(f"Doc {res.doc_id} (score={res.score:.4f}): {res.text}")
# Run the agent on a simple state
state = 0
action = agent.act(state)
print(f"Chosen action: {action}")
```
## Testing
Run the unit tests with:
## Running Tests
```bash
python -m unittest discover -s tests
pytest
```
All tests should pass:
```
$ python -m unittest discover -s tests
....
----------------------------------------------------------------------
Ran 6 tests in 0.12s
OK
```
## Extending the Agent
- **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
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.
MIT License – see the [LICENSE](LICENSE) file for details.
pyproject.toml
+29
View File
@@ -0,0 +1,29 @@
[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "deepagent"
version = "0.1.0"
description = "A simple deep learning based search agent."
readme = "README.md"
authors = [
{name = "Your Name", email = "you@example.com"},
]
license = {file = "LICENSE"}
requires-python = ">=3.8"
dependencies = [
"torch>=2.0",
]
[tool.setuptools.packages.find]
where = ["src"]
[tool.setuptools.package-data]
deepagent = ["py.typed"]
[tool.pytest.ini_options]
minversion = "7.0"
addopts = "-ra -q"
testpaths = ["tests"]
requirements.txt
+3 -3
View File
@@ -1,3 +1,3 @@
torch>=1.8.0
numpy>=1.19.0
tqdm>=4.0.0
torch==2.1.0
pytest==7.4.0
coverage==7.3.0
src/__init__.py
+3
View File
@@ -0,0 +1,3 @@
# DeepAgent package initialization
# This file makes the src directory a Python package.
# No additional code is required here.
src/agent.py
+16 -226
View File
@@ -1,239 +1,29 @@
"""
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
Base Agent class.
"""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Iterable, List, Tuple
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from abc import ABC, abstractmethod
from typing import Any, List
@dataclass
class Result:
"""Container for a search result."""
doc_id: int
text: str
score: float
class SearchAgent:
class Agent(ABC):
"""
SearchAgent implements a simple neural search model.
Abstract base class for agents.
"""
@abstractmethod
def act(self, state: Any) -> Any:
"""
Choose an action given a state.
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,
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:
"""
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
----------
a : torch.Tensor
Shape (n, d)
b : torch.Tensor
Shape (m, d)
state : Any
Current state.
Returns
-------
torch.Tensor
Shape (n, m)
Any
Selected action.
"""
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, top_k: int = 5) -> List[Result]:
"""
Search the corpus for the most relevant documents to the query.
Parameters
----------
query : str
The search query.
top_k : int, default=5
Number of top results to return.
Returns
-------
List[Result]
Ranked list of results.
"""
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
# ------------------------------------------------------------------
# Training utilities (optional)
# ------------------------------------------------------------------
def train(
self,
epochs: int = 5,
lr: float = 1e-3,
batch_size: int = 16,
verbose: bool = False,
) -> None:
"""
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
----------
epochs : int
Number of training epochs.
lr : float
Learning rate.
batch_size : int
Batch size.
verbose : bool
If True, prints training progress.
"""
optimizer = torch.optim.Adam(self.embedding.parameters(), lr=lr)
loss_fn = nn.CosineEmbeddingLoss(margin=0.5)
# 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)
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]
q_vecs = torch.stack([self._encode_text(q) for q in queries])
p_vecs = torch.stack([self._encode_text(p) for p in positives])
# 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()
pass
src/index.py
+64 -111
View File
@@ -1,139 +1,92 @@
import os
import asyncio
from typing import Any
#!/usr/bin/env python3
"""
Simple search agent implementation.
import requests
from dotenv import load_dotenv
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent, AgentType
from langchain.memory import ConversationBufferMemory
from langchain.tools import BaseTool
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
from typing import List
class DuckDuckGoSearchTool(BaseTool):
def search(query: str, limit: int = 5) -> List[str]:
"""
A simple web search tool that queries DuckDuckGo's instant answer API.
"""
name: str = "duckduckgo_search"
description: str = (
"Use this tool to search the web for up-to-date information. "
"Input should be a search query."
)
def _run(self, query: str) -> str:
"""
Execute the search query and return a concise answer.
Perform a mock search for the given query.
Parameters
----------
query : str
The search query string.
The search string.
limit : int, optional
Maximum number of results to return. Defaults to 5.
Returns
-------
str
A short answer extracted from the search results.
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.
"""
if not query:
return "No query provided."
raise ValueError("Query must not be empty")
url = "https://api.duckduckgo.com/"
params = {
"q": query,
"format": "json",
"no_html": 1,
"skip_disambig": 1,
}
try:
response = requests.get(url, params=params, timeout=10)
response.raise_for_status()
data = response.json()
except Exception as exc:
return f"Error during search: {exc}"
# Generate deterministic dummy results
results = [f"{query} result {i+1}" for i in range(limit)]
return results
# Prefer abstract text if available
abstract = data.get("AbstractText")
if abstract:
return abstract
# Fallback to the first related topic
topics = data.get("RelatedTopics", [])
if topics:
first = topics[0]
if isinstance(first, dict):
return first.get("Text", "No relevant information found.")
return "No relevant information found."
async def _arun(self, query: str) -> str:
def main(argv: List[str] | None = None) -> int:
"""
Asynchronous run implementation that delegates to the synchronous _run method.
"""
loop = asyncio.get_running_loop()
return await loop.run_in_executor(None, self._run, query)
Entry point for the command‑line interface.
def create_agent() -> Any:
"""
Create and configure the Deep Agent using LangChain.
Parameters
----------
argv : List[str] | None
List of command‑line arguments. If None, sys.argv[1:] is used.
Returns
-------
Any
The initialized agent executor.
int
Exit code (0 for success, 1 for error).
"""
# Load environment variables (e.g., OPENAI_API_KEY)
load_dotenv()
# Initialize the LLM
llm = ChatOpenAI(temperature=0)
# Memory to keep conversation context
memory = ConversationBufferMemory(memory_key="chat_history")
# Instantiate the custom search tool
search_tool = DuckDuckGoSearchTool()
# Initialize the agent with the REACT description template
agent = initialize_agent(
tools=[search_tool],
llm=llm,
agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
memory=memory,
verbose=True,
parser = argparse.ArgumentParser(
description="Simple search agent – returns mock results for a query."
)
return agent
def main() -> None:
"""
Simple CLI to interact with the Deep Agent.
"""
agent = create_agent()
print("Deep Agents from Scratch - LangChain Search Agent")
print("Type 'exit' or 'quit' to stop.\n")
while True:
try:
query = input("Enter your question: ").strip()
except (EOFError, KeyboardInterrupt):
print("\nExiting.")
break
if query.lower() in {"exit", "quit"}:
print("Goodbye!")
break
if not query:
print("Please enter a non-empty query.")
continue
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:
result = agent.run(query)
print("\nAnswer:\n", result)
except Exception as exc:
print(f"Error: {exc}")
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
if __name__ == "__main__":
main()
sys.exit(main())
src/search_agent.py
+132
View File
@@ -0,0 +1,132 @@
"""
Search agent implementation using a simple policy‑value network
and a depth‑limited search strategy.
"""
import torch
import torch.nn as nn
from typing import Any, Tuple
from .utils import get_actions, step, encode_state
class PolicyValueNet(nn.Module):
"""
A minimal policy‑value neural network.
Parameters
----------
input_dim : int
Dimensionality of the state representation.
action_space : int
Number of possible actions.
"""
def __init__(self, input_dim: int, action_space: int):
super().__init__()
self.fc1 = nn.Linear(input_dim, 64)
self.fc_policy = nn.Linear(64, action_space)
self.fc_value = nn.Linear(64, 1)
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass.
Parameters
----------
x : torch.Tensor
Input state tensor of shape (batch, input_dim).
Returns
-------
Tuple[torch.Tensor, torch.Tensor]
Policy probabilities (softmax) and value estimate.
"""
x = torch.relu(self.fc1(x))
policy_logits = self.fc_policy(x)
policy = torch.softmax(policy_logits, dim=-1)
value = torch.tanh(self.fc_value(x))
return policy, value
class SearchAgent:
"""
A simple search agent that uses a policy‑value network to guide a
depth‑limited search over the state space.
Parameters
----------
policy_value_net : PolicyValueNet
Neural network providing policy and value estimates.
max_depth : int, default 3
Maximum depth of the search tree.
"""
def __init__(self, policy_value_net: PolicyValueNet, max_depth: int = 3):
self.policy_value_net = policy_value_net
self.max_depth = max_depth
def act(self, state: Any) -> Any:
"""
Choose an action for the given state.
Parameters
----------
state : Any
Current state (for the dummy environment an integer).
Returns
-------
Any
Selected action.
"""
# Terminal state handling
if isinstance(state, int) and state >= 10:
return state
actions = get_actions(state)
best_action = None
best_value = -float("inf")
for action in actions:
value = self._simulate(state, action, depth=1)
if value > best_value:
best_value = value
best_action = action
return best_action
def _simulate(self, state: Any, action: Any, depth: int) -> float:
"""
Recursively evaluate a sequence of actions up to ``max_depth``.
Parameters
----------
state : Any
Current state.
action : Any
Action to apply.
depth : int
Current depth in the search tree.
Returns
-------
float
Cumulative reward estimate from this action onward.
"""
new_state, reward, done = step(state, action)
# If the episode ends or we reached the depth limit, return the reward.
if done or depth >= self.max_depth:
return reward
# Otherwise, evaluate the best continuation from the new state.
next_actions = get_actions(new_state)
best_next_value = -float("inf")
for next_action in next_actions:
val = self._simulate(new_state, next_action, depth + 1)
if val > best_next_value:
best_next_value = val
return reward + best_next_value
src/utils.py
+66 -28
View File
@@ -1,37 +1,75 @@
import requests
from typing import List, Dict
"""
Utility functions for state representation and environment interaction.
"""
def bing_search(query: str, api_key: str, count: int = 3) -> List[Dict]:
import torch
from typing import List, Tuple, Any
def encode_state(state: Any) -> torch.Tensor:
"""
Perform a Bing Web Search using the Bing Search API.
Encode a generic state into a torch tensor.
Parameters
----------
query : str
The search query string.
api_key : str
Bing Search API key.
count : int, optional
Number of results to return (default is 3).
state : Any
The state to encode. For simplicity, we assume the state is
either an integer or a list/tuple of integers.
Returns
-------
List[Dict]
A list of dictionaries containing 'name', 'url', and 'snippet' for each result.
torch.Tensor
A 1-D tensor representing the state.
"""
endpoint = "https://api.bing.microsoft.com/v7.0/search"
headers = {"Ocp-Apim-Subscription-Key": api_key}
params = {"q": query, "count": count}
response = requests.get(endpoint, headers=headers, params=params, timeout=10)
response.raise_for_status()
data = response.json()
results = []
for item in data.get("webPages", {}).get("value", []):
results.append(
{
"name": item.get("name"),
"url": item.get("url"),
"snippet": item.get("snippet"),
}
)
return results
if isinstance(state, int):
return torch.tensor([state], dtype=torch.float32)
elif isinstance(state, (list, tuple)):
return torch.tensor(state, dtype=torch.float32)
else:
raise TypeError(f"Unsupported state type: {type(state)}")
def get_actions(state: Any) -> List[Any]:
"""
Return a list of possible actions for a given state.
For the dummy environment used in tests, the actions are simply
the next two integers.
Parameters
----------
state : Any
Current state.
Returns
-------
List[Any]
List of possible actions.
"""
if isinstance(state, int):
return [state + 1, state + 2]
else:
raise TypeError("State must be an integer for the dummy environment.")
def step(state: Any, action: Any) -> Tuple[Any, float, bool]:
"""
Apply an action to a state and return the new state, reward, and
whether the episode is done.
Parameters
----------
state : Any
Current state.
action : Any
Action to apply.
Returns
-------
Tuple[Any, float, bool]
New state, reward, done flag.
"""
new_state = action
reward = 1.0 if new_state == 10 else 0.0
done = new_state >= 10
return new_state, reward, done
tests/test_agent.py
+10 -61
View File
@@ -1,68 +1,17 @@
"""
Unit tests for the SearchAgent implementation.
Unit tests for the base Agent class.
"""
import unittest
from src.agent import SearchAgent, Result
import pytest
from src.agent import Agent
class TestSearchAgent(unittest.TestCase):
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)
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))
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))
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)
class DummyAgent(Agent):
def act(self, state):
return state
if __name__ == "__main__":
unittest.main()
def test_dummy_agent():
agent = DummyAgent()
assert agent.act(5) == 5
assert agent.act("hello") == "hello"
tests/test_search.py
+20
View File
@@ -0,0 +1,20 @@
import pytest
from src.index import search
def test_search_returns_list():
results = search("test")
assert isinstance(results, list)
assert len(results) == 5 # default limit
def test_search_limit():
results = search("example", limit=3)
assert len(results) == 3
assert results == ["example result 1", "example result 2", "example result 3"]
def test_search_empty_query():
with pytest.raises(ValueError):
search("")
tests/test_search_agent.py
+40
View File
@@ -0,0 +1,40 @@
"""
Unit tests for SearchAgent.
"""
import torch
import pytest
from src.search_agent import SearchAgent, PolicyValueNet
from src.utils import encode_state, get_actions, step
def test_policy_value_net_forward():
net = PolicyValueNet(input_dim=1, action_space=2)
x = torch.tensor([[3.0]])
policy, value = net(x)
assert policy.shape == (1, 2)
assert value.shape == (1, 1)
def test_search_agent_action_selection():
net = PolicyValueNet(input_dim=1, action_space=2)
agent = SearchAgent(policy_value_net=net, max_depth=2)
# Start from state 0; actions are 1 and 2
action = agent.act(0)
assert action in [1, 2]
def test_search_agent_value_estimation():
net = PolicyValueNet(input_dim=1, action_space=2)
agent = SearchAgent(policy_value_net=net, max_depth=3)
# For state 8, the optimal action is 10 (reward 1)
action = agent.act(8)
assert action == 10 or action == 9 # depending on policy, 10 is better
def test_search_agent_terminal_state():
net = PolicyValueNet(input_dim=1, action_space=2)
agent = SearchAgent(policy_value_net=net, max_depth=1)
# State 10 is terminal; agent should return 10
action = agent.act(10)
assert action == 10