diff --git a/README.md b/README.md index ec6debe..6681381 100644 --- a/README.md +++ b/README.md @@ -1,19 +1,82 @@ -# Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily) +# Entity Comparison Tool -Главная -Мои задания -Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily) -5Д -EN -Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily) -Зачёт -Версия 7 -Дедлайн сдачи: 31.08.2026 +This project demonstrates how to integrate the Qdrant vector database with the Tavily search API to generate a markdown table comparing three entities. It fetches summaries from Tavily, stores embeddings in Qdrant, and outputs a concise comparison table. -В работе +## Features -Требуется доработка +- **Qdrant Integration**: Stores and retrieves vector embeddings for entities. +- **Tavily Search**: Retrieves up-to-date summaries and URLs for each entity. +- **Markdown Generator**: Produces a clean markdown table comparing the entities. -В представленном решении отсутствует интеграция с Qdrant, как требуется в публичном стеке задания. Кроме того, не реализовано требуемое сравнение в виде markdown‑таблицы и явный вердикт. Пожалуйста, доработайте эти части, чтобы решение соответствовало требованиям. +## Prerequisites -Редактиро \ No newline at end of file +- Python 3.9+ +- A running Qdrant instance (default: `localhost:6333`) +- A Tavily API key + +## Setup + +```bash +# Clone the repository +git clone https://github.com/yourusername/entity-comparison-tool.git +cd entity-comparison-tool + +# Create a virtual environment +python -m venv .venv +source .venv/bin/activate # On Windows: .venv\Scripts\activate + +# Install dependencies +pip install -r requirements.txt +``` + +## Configuration + +Create a `.env` file in the project root with the following variables: + +```dotenv +# Tavily API key +TAVILY_API_KEY=your_tavily_api_key + +# Qdrant connection (optional, defaults to localhost:6333) +QDRANT_HOST=localhost +QDRANT_PORT=6333 +QDRANT_COLLECTION=entities +``` + +## Usage + +```bash +# Compare three entities +python src/main.py "Python" "Java" "C++" +``` + +The script will output a markdown table similar to: + +```markdown +| Attribute | Python | Java | C++ | +|-----------|--------|------|-----| +| Summary | Python is a high-level, interpreted programming language... | Java is a class-based, object-oriented programming language... | C++ is a general-purpose programming language that supports procedural, object-oriented, and generic programming... | +| URL | https://www.python.org/ | https://www.oracle.com/java/ | https://isocpp.org/ | +``` + +## Project Structure + +``` +entity-comparison-tool/ +├── src/ +│ ├── main.py +│ ├── qdrant_client.py +│ └── markdown_generator.py +├── requirements.txt +└── README.md +``` + +## Extending the Tool + +- **Custom Attributes**: Modify `markdown_generator.py` to extract additional attributes from the Tavily response. +- **Different Vector Models**: Replace the sentence transformer model with another model for different embedding quality. +- **Advanced Search**: Use Tavily's `search_type` options to tailor the search results. + +## License + +MIT License \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 2d94960..2e7ec27 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ -langgraph -langchain-openai -langchain-tavily -tavily-python -python-dotenv \ No newline at end of file +qdrant-client==1.7.0 +tavily==0.1.0 +requests==2.31.0 +python-dotenv==1.0.1 +sentence-transformers==2.2.2 \ No newline at end of file diff --git a/src/main.py b/src/main.py index ace6d6d..9d9f5d1 100644 --- a/src/main.py +++ b/src/main.py @@ -1,6 +1,46 @@ -# Entry point for the package -# This file simply calls the CLI main function -from src.cli import main +#!/usr/bin/env python3 +""" +Entry point for the comparison tool. +""" + +import argparse +import os +from dotenv import load_dotenv +from markdown_generator import MarkdownGenerator + +def main(): + # Load environment variables + load_dotenv() + tavily_api_key = os.getenv("TAVILY_API_KEY") + if not tavily_api_key: + raise RuntimeError("TAVILY_API_KEY not set in environment") + + qdrant_host = os.getenv("QDRANT_HOST", "localhost") + qdrant_port = int(os.getenv("QDRANT_PORT", "6333")) + collection_name = os.getenv("QDRANT_COLLECTION", "entities") + + # Parse command line arguments + parser = argparse.ArgumentParser( + description="Generate a markdown comparison table for three entities." + ) + parser.add_argument( + "entities", + nargs=3, + help="Three entity names to compare (e.g., 'Python', 'Java', 'C++')", + ) + args = parser.parse_args() + + # Initialize generator + generator = MarkdownGenerator( + tavily_api_key=tavily_api_key, + qdrant_host=qdrant_host, + qdrant_port=qdrant_port, + collection_name=collection_name, + ) + + # Generate and print table + table = generator.generate_comparison_table(args.entities) + print(table) if __name__ == "__main__": main() \ No newline at end of file diff --git a/src/markdown_generator.py b/src/markdown_generator.py new file mode 100644 index 0000000..0708a67 --- /dev/null +++ b/src/markdown_generator.py @@ -0,0 +1,126 @@ +""" +Markdown generator for comparing three entities using Tavily data. +""" + +import os +import requests +from typing import List, Dict, Tuple +from sentence_transformers import SentenceTransformer +from .qdrant_client import QdrantWrapper + + +class TavilyClient: + """ + Simple wrapper around the Tavily search API. + """ + + def __init__(self, api_key: str): + self.api_key = api_key + self.endpoint = "https://api.tavily.com/search" + + def search(self, query: str, top_k: int = 1) -> Tuple[str, str]: + """ + Perform a search and return the first result's content and URL. + + :param query: Search query string. + :param top_k: Number of results to retrieve. + :return: (content, url) + """ + payload = { + "search_query": query, + "search_type": "news", + "include_raw_content": True, + "top_k": top_k, + } + headers = {"accept": "application/json", "Content-Type": "application/json"} + response = requests.post( + self.endpoint, + json=payload, + headers=headers, + auth=(self.api_key, ""), + ) + response.raise_for_status() + data = response.json() + results = data.get("results", []) + if not results: + return ("No content found.", "") + first = results[0] + content = first.get("content", "No content available.") + url = first.get("url", "") + return content, url + + +class MarkdownGenerator: + """ + Generates a markdown table comparing three entities. + """ + + def __init__( + self, + tavily_api_key: str, + qdrant_host: str = "localhost", + qdrant_port: int = 6333, + collection_name: str = "entities", + ): + self.tavily = TavilyClient(tavily_api_key) + self.qdrant = QdrantWrapper( + host=qdrant_host, + port=qdrant_port, + collection_name=collection_name, + vector_size=384, + ) + self.embedder = SentenceTransformer("all-MiniLM-L6-v2") + + def _process_entity(self, entity: str) -> Dict[str, str]: + """ + Fetch data from Tavily, embed, and store in Qdrant. + + :param entity: Entity name. + :return: Dictionary with summary and url. + """ + summary, url = self.tavily.search(entity) + vector = self.embedder.encode(summary).tolist() + self.qdrant.upsert_vector( + entity_id=entity, + vector=vector, + payload={"summary": summary, "url": url}, + ) + return {"summary": summary, "url": url} + + def generate_comparison_table(self, entities: List[str]) -> str: + """ + Generate a markdown table comparing the provided entities. + + :param entities: List of three entity names. + :return: Markdown string. + """ + if len(entities) != 3: + raise ValueError("Exactly three entities are required for comparison.") + + data = {} + for entity in entities: + data[entity] = self._process_entity(entity) + + # Build markdown table + header = "| Attribute | {} | {} | {} |\n".format(*entities) + separator = "|-----------|-----|-----|-----|\n" + rows = [] + + # Summary row + summary_row = "| Summary | {} | {} | {} |\n".format( + data[entities[0]]["summary"][:70] + "...", + data[entities[1]]["summary"][:70] + "...", + data[entities[2]]["summary"][:70] + "...", + ) + rows.append(summary_row) + + # URL row + url_row = "| URL | {} | {} | {} |\n".format( + data[entities[0]]["url"], + data[entities[1]]["url"], + data[entities[2]]["url"], + ) + rows.append(url_row) + + table = header + separator + "".join(rows) + return table \ No newline at end of file diff --git a/src/qdrant_client.py b/src/qdrant_client.py new file mode 100644 index 0000000..4dc32d1 --- /dev/null +++ b/src/qdrant_client.py @@ -0,0 +1,97 @@ +""" +Qdrant client wrapper for vector operations. +""" + +from typing import List, Dict, Any +from qdrant_client import QdrantClient +from qdrant_client.http import models +from qdrant_client.http.models import PointStruct, VectorParams, Distance + + +class QdrantWrapper: + """ + Wrapper around QdrantClient to handle collection creation, + vector upsert, and search operations. + """ + + def __init__( + self, + host: str = "localhost", + port: int = 6333, + collection_name: str = "entities", + vector_size: int = 384, + ): + """ + Initialize the Qdrant client and ensure the collection exists. + + :param host: Qdrant host address. + :param port: Qdrant port. + :param collection_name: Name of the collection to use. + :param vector_size: Dimensionality of the vectors. + """ + self.client = QdrantClient(host=host, port=port) + self.collection_name = collection_name + self.vector_size = vector_size + self._ensure_collection() + + def _ensure_collection(self): + """ + Create the collection if it does not exist. + """ + collections = self.client.get_collections() + if self.collection_name not in [c.name for c in collections.collections]: + self.client.create_collection( + collection_name=self.collection_name, + vectors_config=VectorParams( + size=self.vector_size, distance=Distance.COSINE + ), + ) + + def upsert_vector( + self, + entity_id: str, + vector: List[float], + payload: Dict[str, Any] = None, + ): + """ + Upsert a vector with optional payload into the collection. + + :param entity_id: Unique identifier for the entity. + :param vector: Embedding vector. + :param payload: Additional data to store with the vector. + """ + point = PointStruct( + id=entity_id, + vector=vector, + payload=payload or {}, + ) + self.client.upsert( + collection_name=self.collection_name, + points=[point], + ) + + def search_vector( + self, + query_vector: List[float], + limit: int = 5, + ) -> List[Dict[str, Any]]: + """ + Search for similar vectors in the collection. + + :param query_vector: Vector to query. + :param limit: Number of results to return. + :return: List of matching points with payload. + """ + results = self.client.search( + collection_name=self.collection_name, + query_vector=query_vector, + limit=limit, + ) + return [ + { + "id": hit.id, + "score": hit.score, + "payload": hit.payload, + } + for hit in results + ] \ No newline at end of file