src/database/core.py - master - gurgutan/tutor

src/database/core.py
183 строки6 KB
Ivan Slepovichev
+llms.py, +async_assistant.py
12 фев 2025, 16:11
12 фев 2025, 16:11feccf21
from __future__ import annotations
from datetime import datetime
from pathlib import Path
from typing import List, Union, Dict, Any

import yaml
from txtai.embeddings import Embeddings

from .models import Document
from src.utils.loggers import logger, logit


class VectorDatabase:
    """Vector database management class using txtai.

    This class provides methods for managing document collections
    in a vector database with similarity search capabilities.

    Example:
        db = VectorDatabase("config.yml")
        docs = [Document(id="1", content="Sample text", metadata={})]
        db.add_documents(docs)
        results = db.search("query text", limit=5)
    """

    def __init__(self, config_path: Union[str, Path]) -> None:
        """Initialize vector database with configuration.

        Args:
            config_path: Path to YAML configuration file

        Example:
            db = VectorDatabase("config.yml")
        """
        self.config_path: Path = Path(config_path)
        self.config: Dict[str, Any] = self._load_config()
        self.embeddings: Embeddings = Embeddings(config=self.config["embeddings"])
        logger.debug(f"VectorDatabase initialized with config from {self.config_path}")
        self._load_index()

    def index(self, documents: List[Document]) -> VectorDatabase:
        """Initialize database and create necessary structures.

        Args:
            documents: List of Document objects

        Example:
            db.index(docs)
        """
        data = [(doc.id, {"text": doc.content, **doc.metadata}) for doc in documents]
        self.embeddings.index(data)
        self._save_index()
        logger.debug(f"Documents indexed: {[doc.id for doc in documents]}")
        return self

    def add_documents(self, documents: List[Document]) -> VectorDatabase:
        """Add collection of documents to database.

        Args:
            documents: List of Document objects to add

        Example:
            db.add_documents(docs)
        """
        data = [(doc.id, {"text": doc.content, **doc.metadata}) for doc in documents]
        for doc in data:
            self.embeddings.upsert([doc])
        self._save_index()
        logger.debug(f"Documents indexed: {[doc.id for doc in documents]}")
        return self

    def delete_documents(self, doc_ids: List[str]) -> VectorDatabase:
        """Delete documents by their IDs.

        Args:
            doc_ids: List of document IDs to delete

        Example:
            db.delete_documents(["1", "2"])
        """
        self.embeddings.delete(doc_ids)
        logger.debug(f"Documents deleted: {doc_ids}")
        return self

    def update_document(self, doc: Document) -> VectorDatabase:
        """Update existing document content.

        Args:
            document: Document object with updated content

        Example:
            db.update_document(doc)
        """
        doc.updated_at = datetime.now()
        self.embeddings.upsert([(doc.id, {"text": doc.content, **doc.metadata})])
        self._save_index()
        return self

    def search(self, query: str, limit: int = 2) -> List[Dict[str, Any]]:
        """Search documents by semantic similarity.

        Args:
            query: Search query text
            limit: Maximum number of results to return

        Returns:
            List of documents sorted by relevance

        Example:
            results = db.search("query", limit=5)
        """
        results = self.embeddings.search(query, limit)
        logger.debug(f"Search performed with query: {query}, limit: {limit}")
        return results

    def get_all(self, limit: int = 64) -> List[Dict[str, Any]]:
        """Get all documents from the database.

        Args:
            limit: Maximum number of documents to return

        Returns:
            List of documents sorted by relevance

        Example:
            results = db.get_all(limit=5)
        """
        query = "SELECT id, text, title FROM txtai"
        results = self.embeddings.search(query, limit=limit)
        return results

    def _load_config(self) -> Dict[str, Any]:
        """Load database configuration from YAML file.

        Returns:
            Dictionary with configuration parameters

        Example:
            config = self._load_config()
        """
        if not self.config_path.exists():
            raise FileNotFoundError(f"Config file not found: {self.config_path}")
        with open(self.config_path) as f:
            config_data = yaml.safe_load(f)
        logger.debug(f"Configuration loaded: {config_data}")
        return config_data

    def _load_index(self) -> None:
        """Load the vector index from the defined path in configuration.

        Example:
            self._load_index()
        """
        index_path = self.config.get("path")
        if not index_path:
            logger.debug(
                "No index path specified in configuration; skipping loading index."
            )
            return
        index_file: Path = Path(index_path)
        if index_file.exists():
            self.embeddings.load(str(index_file))
            logger.debug(f"Index loaded from {index_file}")
        else:
            logger.debug(
                f"Index file {index_file} does not exist; skipping loading index."
            )

    def _save_index(self) -> None:
        """Save the vector index to the defined path in configuration.

        Example:
            self._save_index()
        """
        index_path = self.config.get("path")
        if not index_path:
            logger.debug(
                "No index path specified in configuration; skipping saving index."
            )
            return
        index_file: Path = Path(index_path)
        self.embeddings.save(str(index_file))
        logger.debug(f"Index saved to {index_file}")