ragit 0.1__tar.gz

ragit-0.1/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.2
+Name: ragit
+Version: 0.1
+Requires-Dist: sentence-transformers>=3.4.1
+Requires-Dist: pandas>=2.2.3
+Requires-Dist: chromadb>=0.6.3
+Requires-Dist: setuptools>=75.8.0
+Requires-Dist: wheel>=0.45.1
+Requires-Dist: twine>=6.1.0
+Dynamic: requires-dist

ragit-0.1/README.md

@@ -0,0 +1,153 @@
+# ragit
+
+**ragit** is a lightweight Python library that simplifies the management of vector databases using [ChromaDB](https://github.com/chroma-core/chroma) and [Sentence Transformers](https://www.sbert.net/). With ragit, you can easily create, update, query, and manage your vector database—all from CSV files containing text data.
+
+## Features
+
+- **Create a Vector Database:** Build your database from a CSV file with two required columns: `id` and `text`.
+- **Add New Entries:** Insert additional entries from CSV files or add them individually.
+- **Similarity Search:** Find nearby texts using various distance metrics (e.g., cosine, L2) with similarity scores.
+- **Data Retrieval:** Fetch entries by IDs or exact text matches.
+- **Deletion:** Remove single entries or entire collections when needed.
+
+## Installation
+
+Install ragit from PyPI using pip:
+
+```bash
+pip install ragit
+```
+
+## CSV File Format
+ragit expects your CSV file to have exactly two columns: `id` and `text`.
+
+## Example CSV (`data.csv`):
+
+```csv
+id,text
+1,The quick brown fox jumps over the lazy dog.
+2,Another sample entry for testing.
+```
+
+## Usage
+Below are some examples that demonstrate how to use `ragit`. The examples cover creating a database, adding entries, performing similarity searches, and more.
+
+### 1. Importing and Initializing
+First, import the `VectorDBManager` class from `ragit` and initialize it:
+
+```python
+from ragit import VectorDBManager
+
+# Initialize the vector database manager with a custom persistence directory and model
+db_manager = VectorDBManager(
+    persist_directory="./my_vector_db", # Optional # default : "./vector_db"
+    provider="sentence_transformer", # Optional # default : "sentence_transformer"
+    model_name="all-mpnet-base-v2" # Optional # default : "all-mpnet-base-v2"
+)
+```
+
+### 2. Creating a Database
+Create a new collection (named `my_collection`) using your CSV file. In this example, the `distance_metric` is set to "cosine"(available options: l2, cosine, ip, l1) :
+
+```python
+db_manager.create_database(
+    csv_path="data.csv", 
+    collection_name="my_collection",
+    distance_metric="cosine" # Optional # default : l2
+)
+```
+
+### 3. Adding a Single Entry
+Add an individual entry to the collection:
+
+```python
+db_manager.add_single_row(
+    id_="101",
+    text="This is a new test entry for the database.",
+    collection_name="my_collection"
+)
+```
+
+### 4. Adding Multiple Entries from CSV
+You can also add multiple entries from a CSV file. This function skips any entries that already exist in the collection:
+
+```python
+stats = db_manager.add_values_from_csv(
+    csv_path="data.csv",
+    collection_name="my_collection"
+)
+print(f"Added {stats['new_entries_added']} new entries")
+```
+
+### 5. Retrieving Collection Information
+Fetch and display information about your collection:
+
+```python
+info = db_manager.get_collection_info("my_collection")
+print(f"Collection size: {info['count']} entries")
+```
+
+### 6. Performing a Similarity Search
+Find texts that are similar to your query. In this example, the query text is "ai", and the search is filtered using the string "Artificial intelligence". The top 2 results are returned:
+
+```python
+results = db_manager.find_nearby_texts(
+    text="ai",
+    collection_name="my_collection",
+    k=2,
+    search_string="Artificial intelligence" # Optional
+)
+
+print("Results:")
+for item in results:
+    print(f"\nID: {item['id']}")
+    print(f"Text: {item['text']}")
+    print(f"Similarity: {item['similarity']}%")
+    print(f"Distance ({item['metric']}): {item['raw_distance']}")
+```
+
+### 7. Deleting an Entry
+Remove an entry from the collection by its ID:
+
+```python
+db_manager.delete_entry_by_id(
+    id_="1",
+    collection_name="my_collection"
+)
+```
+
+### 8. Fetching Texts by IDs
+Retrieve text entries for a list of IDs:
+
+```python
+ids_to_fetch = ["1", "2", "3"]
+texts = db_manager.get_by_ids(ids_to_fetch, "my_collection")
+print("Texts:", texts)
+```
+
+### 9. Fetching IDs by Texts
+For an exact text match, get the corresponding IDs:
+
+```python
+texts_to_fetch = [
+    "Plato was an ancient Greek philosopher of the Classical period who is considered a foundational thinker in Western philosophy"
+]
+ids = db_manager.get_by_texts(texts_to_fetch, "my_collection")
+print("IDs:", ids)
+```
+
+### 10. Deleting a Collection
+Delete an entire collection. **Note:** You must pass `confirmation="yes"` to proceed with deletion.
+
+```python
+db_manager.delete_collection(
+    collection_name="my_collection",
+    confirmation="yes"
+)
+```
+
+## Contributing
+Contributions are welcome! If you encounter any issues or have suggestions for improvements, please feel free to open an issue or submit a pull request on GitHub.
+
+## License
+This project is licensed under the MIT License. See the `LICENSE` file for details.

ragit-0.1/ragit/__init__.py

@@ -0,0 +1,2 @@
+# __init__.py
+from .main import VectorDBManager

ragit-0.1/ragit/main.py

@@ -0,0 +1,384 @@
+import chromadb
+import pandas as pd
+import logging
+from sentence_transformers import SentenceTransformer
+from typing import List, Dict, Optional, Union
+import os
+
+
+class VectorDBManager:
+    def __init__(
+        self,
+        persist_directory: str = "./vector_db",
+        provider: str = "sentence_transformer",
+        model_name: str = "all-mpnet-base-v2",
+    ):
+        """
+        Initialize the Vector Database Manager.
+
+        Args:
+            persist_directory (str): Directory to persist the database
+        """
+        self.persist_directory = persist_directory
+        self.client = chromadb.PersistentClient(path=persist_directory)
+        if provider == "sentence_transformer":
+            self.model = SentenceTransformer(model_name)
+
+        logging.basicConfig(
+            level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+        )
+        self.logger = logging.getLogger(__name__)
+
+    def create_database(
+        self,
+        csv_path: str,
+        collection_name: str,
+        distance_metric: str = "l2",
+        collection_metadata: Dict = None,
+    ) -> bool:
+        """
+        Create a new database from a CSV file.
+
+        Args:
+            csv_path (str): Path to the CSV file containing 'id' and 'text' columns
+            collection_name (str): Name of the collection to create
+            distance_metric (str): Distance metric (l2, cosine, ip)
+            collection_metadata (Dict, optional): Additional metadata for the collection
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+
+            df = pd.read_csv(csv_path)
+
+            if not {"id", "text"}.issubset(df.columns):
+                self.logger.error("CSV must contain 'id' and 'text' columns")
+                return False
+
+            collection_meta = {
+                "hnsw:space": distance_metric,
+                "description": f"Collection created from {csv_path}",
+            }
+
+            if collection_metadata:
+                collection_meta.update(collection_metadata)
+
+            collection = self.client.create_collection(
+                name=collection_name, metadata=collection_meta
+            )
+
+            embeddings = self.model.encode(df["text"].tolist()).tolist()
+
+            collection.add(
+                ids=[str(id_) for id_ in df["id"]],
+                documents=df["text"].tolist(),
+                embeddings=embeddings,
+            )
+
+            self.logger.info(f"Successfully created collection '{collection_name}'")
+            return True
+
+        except Exception as e:
+            self.logger.error(f"Error creating database: {str(e)}")
+            return False
+
+    def add_values_from_csv(
+        self, csv_path: str, collection_name: str
+    ) -> Dict[str, int]:
+        """
+        Add values from CSV file to existing collection, skipping existing IDs.
+
+        Args:
+            csv_path (str): Path to the CSV file
+            collection_name (str): Name of the target collection
+
+        Returns:
+            Dict[str, int]: Statistics about the operation
+        """
+        try:
+
+            df = pd.read_csv(csv_path)
+
+            collection = self.client.get_collection(collection_name)
+
+            existing_ids = set(collection.get()["ids"])
+
+            new_df = df[~df["id"].astype(str).isin(existing_ids)]
+
+            if not new_df.empty:
+
+                embeddings = self.model.encode(new_df["text"].tolist()).tolist()
+
+                collection.add(
+                    ids=[str(id_) for id_ in new_df["id"]],
+                    documents=new_df["text"].tolist(),
+                    embeddings=embeddings,
+                )
+
+            stats = {
+                "total_entries": len(df),
+                "new_entries_added": len(new_df),
+                "skipped_entries": len(df) - len(new_df),
+            }
+
+            self.logger.info(
+                f"Added {stats['new_entries_added']} new entries to '{collection_name}'"
+            )
+            return stats
+
+        except Exception as e:
+            self.logger.error(f"Error adding values from CSV: {str(e)}")
+            return {"error": str(e)}
+
+    def add_single_row(self, id_: str, text: str, collection_name: str) -> bool:
+        """
+        Add a single entry to the collection.
+
+        Args:
+            id_ (str): ID for the new entry
+            text (str): Text content
+            collection_name (str): Target collection name
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+            collection = self.client.get_collection(collection_name)
+
+            if str(id_) in collection.get()["ids"]:
+                self.logger.warning(f"ID {id_} already exists in collection")
+                return False
+
+            embedding = self.model.encode([text]).tolist()
+
+            collection.add(ids=[str(id_)], documents=[text], embeddings=embedding)
+
+            self.logger.info(f"Successfully added entry with ID {id_}")
+            return True
+
+        except Exception as e:
+            self.logger.error(f"Error adding single row: {str(e)}")
+            return False
+
+    def delete_entry_by_id(self, id_: str, collection_name: str) -> bool:
+        """
+        Delete an entry by its ID.
+
+        Args:
+            id_ (str): ID of the entry to delete
+            collection_name (str): Collection name
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+            collection = self.client.get_collection(collection_name)
+
+            if str(id_) not in collection.get()["ids"]:
+                self.logger.warning(f"ID {id_} not found in collection")
+                return False
+
+            collection.delete(ids=[str(id_)])
+
+            self.logger.info(f"Successfully deleted entry with ID {id_}")
+            return True
+
+        except Exception as e:
+            self.logger.error(f"Error deleting entry: {str(e)}")
+            return False
+
+    def find_nearby_texts(
+        self,
+        text: str,
+        collection_name: str,
+        search_string: Optional[str] = None,
+        k: int = 5,
+    ) -> List[Dict[str, Union[str, float]]]:
+        """
+        Find nearby texts using similarity search with scores.
+
+        Args:
+            text (str): Query text
+            collection_name (str): Collection to search in
+            k (int): Number of results to return
+
+        Returns:
+            List[Dict[str, Union[str, float]]]: List of nearby texts with their IDs and similarity scores
+        """
+        try:
+            collection = self.client.get_collection(collection_name)
+            print("Metadata:", collection.metadata)
+
+            distance_metric = collection.metadata["hnsw:space"]
+
+            query_embedding = self.model.encode([text]).tolist()
+
+            if search_string:
+                results = collection.query(
+                    query_embeddings=query_embedding,
+                    n_results=k,
+                    include=["documents", "distances", "metadatas"],
+                    where_document={"$contains": search_string},
+                )
+            else:
+                results = collection.query(
+                    query_embeddings=query_embedding,
+                    n_results=k,
+                    include=["documents", "distances", "metadatas"],
+                )
+
+            distances = results["distances"][0]
+            if not distances:
+                return []
+
+            similarities = []
+            for dist in distances:
+                if distance_metric == "cosine":
+
+                    similarity = 1 - dist
+                elif distance_metric == "ip":
+
+                    min_dist = min(distances)
+                    max_dist = max(distances)
+                    similarity = (
+                        (dist - min_dist) / (max_dist - min_dist)
+                        if max_dist > min_dist
+                        else 1.0
+                    )
+                elif distance_metric == "l1":
+
+                    max_dist = max(distances)
+                    similarity = 1 - (dist / max_dist) if max_dist > 0 else 1.0
+                elif distance_metric == "l2":
+
+                    max_dist = max(distances)
+                    similarity = 1 - (dist / max_dist) if max_dist > 0 else 1.0
+
+                similarities.append(similarity)
+
+            nearby_texts = [
+                {
+                    "id": id_,
+                    "text": text_,
+                    "similarity": round(similarity * 100, 4),
+                    "raw_distance": dist,
+                    "metric": distance_metric,
+                }
+                for id_, text_, similarity, dist in zip(
+                    results["ids"][0], results["documents"][0], similarities, distances
+                )
+            ]
+
+            return nearby_texts
+
+        except Exception as e:
+            self.logger.error(f"Error finding nearby texts: {str(e)}")
+            return []
+
+    def delete_collection(self, collection_name: str, confirmation: str = "no") -> bool:
+        """
+        Delete an entire collection.
+
+        Args:
+            collection_name (str): Name of collection to delete
+            confirmation (str): Must be 'yes' to proceed
+
+        Returns:
+            bool: True if successful, False otherwise
+        """
+        try:
+            if confirmation.lower() != "yes":
+                self.logger.warning("Deletion cancelled - confirmation not provided")
+                return False
+
+            self.client.delete_collection(collection_name)
+            self.logger.info(f"Successfully deleted collection '{collection_name}'")
+            return True
+
+        except Exception as e:
+            self.logger.error(f"Error deleting collection: {str(e)}")
+            return False
+
+    def get_collection_info(self, collection_name: str) -> Dict:
+        """
+        Get information about a collection.
+
+        Args:
+            collection_name (str): Name of the collection
+
+        Returns:
+            Dict: Collection information and statistics
+        """
+        try:
+            collection = self.client.get_collection(collection_name)
+            collection_data = collection.get()
+
+            info = {
+                "name": collection_name,
+                "count": len(collection_data["ids"]),
+                "metadata": collection.metadata,
+            }
+
+            return info
+
+        except Exception as e:
+            self.logger.error(f"Error getting collection info: {str(e)}")
+            return {"error": str(e)}
+
+    def get_by_ids(self, ids: List[str], collection_name: str) -> Dict[str, str]:
+        """
+        Get texts for given IDs in batch.
+
+        Args:
+            ids (List[str]): List of IDs to fetch
+            collection_name (str): Name of the collection
+
+        Returns:
+            Dict[str, str]: Dictionary mapping IDs to their corresponding texts
+        """
+        try:
+            collection = self.client.get_collection(collection_name)
+
+            str_ids = [str(id_) for id_ in ids]
+
+            results = collection.get(ids=str_ids, include=["documents"])
+
+            id_to_text = {
+                id_: text for id_, text in zip(results["ids"], results["documents"])
+            }
+
+            return id_to_text
+
+        except Exception as e:
+            self.logger.error(f"Error getting texts by IDs: {str(e)}")
+            return {}
+
+    def get_by_texts(self, texts: List[str], collection_name: str) -> Dict[str, str]:
+        """
+        Get IDs for given texts in batch.
+        Note: For exact text matching. For similar texts, use find_nearby_texts.
+
+        Args:
+            texts (List[str]): List of texts to fetch
+            collection_name (str): Name of the collection
+
+        Returns:
+            Dict[str, str]: Dictionary mapping texts to their corresponding IDs
+        """
+        try:
+            collection = self.client.get_collection(collection_name)
+
+            all_data = collection.get()
+
+            text_to_id = {
+                text: id_
+                for text, id_ in zip(all_data["documents"], all_data["ids"])
+                if text in texts
+            }
+
+            return text_to_id
+
+        except Exception as e:
+            self.logger.error(f"Error getting IDs by texts: {str(e)}")
+            return {}

ragit-0.1/ragit.egg-info/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.2
+Name: ragit
+Version: 0.1
+Requires-Dist: sentence-transformers>=3.4.1
+Requires-Dist: pandas>=2.2.3
+Requires-Dist: chromadb>=0.6.3
+Requires-Dist: setuptools>=75.8.0
+Requires-Dist: wheel>=0.45.1
+Requires-Dist: twine>=6.1.0
+Dynamic: requires-dist

ragit-0.1/ragit.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+README.md
+setup.py
+ragit/__init__.py
+ragit/main.py
+ragit.egg-info/PKG-INFO
+ragit.egg-info/SOURCES.txt
+ragit.egg-info/dependency_links.txt
+ragit.egg-info/requires.txt
+ragit.egg-info/top_level.txt

ragit-0.1/ragit.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

ragit-0.1/ragit.egg-info/requires.txt

@@ -0,0 +1,6 @@
+sentence-transformers>=3.4.1
+pandas>=2.2.3
+chromadb>=0.6.3
+setuptools>=75.8.0
+wheel>=0.45.1
+twine>=6.1.0

ragit-0.1/ragit.egg-info/top_level.txt

@@ -0,0 +1 @@
+ragit

ragit-0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

ragit-0.1/setup.py

@@ -0,0 +1,11 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="ragit",
+    version="0.1",
+    packages=find_packages(),
+    install_requires = ['sentence-transformers>=3.4.1', 
+                        'pandas>=2.2.3', 'chromadb>=0.6.3', 
+                        'setuptools>=75.8.0', 
+                        'wheel>=0.45.1', 'twine>=6.1.0']   
+)