feat: solution for 'Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily)'

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()

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

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
+        ]