cognee-community-vector-adapter-redis 0.0.3__tar.gz → 0.1.0__tar.gz

{cognee_community_vector_adapter_redis-0.0.3 → cognee_community_vector_adapter_redis-0.1.0}/PKG-INFO

@@ -1,10 +1,12 @@
 Metadata-Version: 2.4
 Name: cognee-community-vector-adapter-redis
-Version: 0.0.3
+Version: 0.1.0
 Summary: Redis vector database adapter for cognee
 Requires-Python: <=3.13,>=3.11
-Requires-Dist: cognee>=0.2.0.dev0
+Requires-Dist: cognee==0.5.2
+Requires-Dist: instructor>=1.11
 Requires-Dist: redisvl<=1.0.0,>=0.6.0
+Requires-Dist: starlette>=0.48.0
 Description-Content-Type: text/markdown
 
 <div align="center" dir="auto">

{cognee_community_vector_adapter_redis-0.0.3 → cognee_community_vector_adapter_redis-0.1.0}/cognee_community_vector_adapter_redis/redis_adapter.py

@@ -1,23 +1,22 @@
-import json
 import asyncio
-from typing import Dict, List, Optional, Any
+import json
+from typing import Any
 from uuid import UUID
 
-from redisvl.index import AsyncSearchIndex
-from redisvl.schema import IndexSchema
-from redisvl.query import VectorQuery
-# from redisvl.query import VectorDistanceMetric
-
-from cognee.shared.logging_utils import get_logger
-
-from cognee.infrastructure.engine import DataPoint
-from cognee.infrastructure.engine.utils import parse_id
 from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
 from cognee.infrastructure.databases.vector import VectorDBInterface
-from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import (
     EmbeddingEngine,
 )
+from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult
+from cognee.infrastructure.engine import DataPoint
+from cognee.infrastructure.engine.utils import parse_id
+
+# from redisvl.query import VectorDistanceMetric
+from cognee.shared.logging_utils import get_logger
+from redisvl.index import AsyncSearchIndex
+from redisvl.query import VectorQuery
+from redisvl.schema import IndexSchema
 
 logger = get_logger("RedisAdapter")
 
@@ -73,15 +72,16 @@ class RedisAdapter(VectorDBInterface):
     """
 
     name = "Redis"
-    url: Optional[str]
-    api_key: Optional[str] = None
-    embedding_engine: Optional[EmbeddingEngine] = None
+    url: str | None
+    api_key: str | None = None
+    embedding_engine: EmbeddingEngine | None = None
 
     def __init__(
         self,
         url: str,
-        api_key: Optional[str] = None,
-        embedding_engine: Optional[EmbeddingEngine] = None,
+        database_name: str = "cognee",
+        api_key: str | None = None,
+        embedding_engine: EmbeddingEngine | None = None,
     ) -> None:
         """Initialize the Redis adapter.
 
@@ -99,11 +99,12 @@ class RedisAdapter(VectorDBInterface):
             raise VectorEngineInitializationError("Embedding engine is required!")
 
         self.url = url
+        self.database_name = database_name
         self.embedding_engine = embedding_engine
         self._indices = {}
         self.VECTOR_DB_LOCK = asyncio.Lock()
 
-    async def embed_data(self, data: List[str]) -> List[List[float]]:
+    async def embed_data(self, data: list[str]) -> list[list[float]]:
         """Embed text data using the embedding engine.
 
         Args:
@@ -179,7 +180,6 @@ class RedisAdapter(VectorDBInterface):
         try:
             index = self._get_index(collection_name)
             result = await index.exists()
-            await index.disconnect()
             return result
         except Exception:
             return False
@@ -187,7 +187,7 @@ class RedisAdapter(VectorDBInterface):
     async def create_collection(
         self,
         collection_name: str,
-        payload_schema: Optional[Any] = None,
+        payload_schema: Any | None = None,
     ) -> None:
         """Create a new collection (Redis index) with vector search capabilities.
 
@@ -205,7 +205,6 @@ class RedisAdapter(VectorDBInterface):
                     logger.info(f"Collection {collection_name} already exists")
                     return
 
-                index = self._get_index(collection_name)
                 await index.create(overwrite=False)
 
                 logger.info(f"Created collection {collection_name}")
@@ -213,12 +212,8 @@ class RedisAdapter(VectorDBInterface):
             except Exception as e:
                 logger.error(f"Error creating collection {collection_name}: {str(e)}")
                 raise e
-            finally:
-                await index.disconnect()
 
-    async def create_data_points(
-        self, collection_name: str, data_points: List[DataPoint]
-    ) -> None:
+    async def create_data_points(self, collection_name: str, data_points: list[DataPoint]) -> None:
         """Create data points in the collection.
 
         Args:
@@ -232,21 +227,16 @@ class RedisAdapter(VectorDBInterface):
         index = self._get_index(collection_name)
         try:
             if not await self.has_collection(collection_name):
-                raise CollectionNotFoundError(
-                    f"Collection {collection_name} not found!"
-                )
+                raise CollectionNotFoundError(f"Collection {collection_name} not found!")
 
             # Embed the data points
             data_vectors = await self.embed_data(
-                [
-                    DataPoint.get_embeddable_data(data_point)
-                    for data_point in data_points
-                ]
+                [DataPoint.get_embeddable_data(data_point) for data_point in data_points]
             )
 
             # Prepare documents for RedisVL
             documents = []
-            for data_point, embedding in zip(data_points, data_vectors):
+            for data_point, embedding in zip(data_points, data_vectors, strict=False):
                 # Serialize the payload to handle UUIDs and other non-JSON types
                 payload = serialize_for_json(data_point.model_dump())
 
@@ -265,19 +255,13 @@ class RedisAdapter(VectorDBInterface):
             # Load using RedisVL
             await index.load(documents, id_field="id")
 
-            logger.info(
-                f"Created {len(data_points)} data points in collection {collection_name}"
-            )
+            logger.info(f"Created {len(data_points)} data points in collection {collection_name}")
 
         except Exception as e:
             logger.error(f"Error creating data points: {str(e)}")
             raise e
-        finally:
-            await index.disconnect()
 
-    async def create_vector_index(
-        self, index_name: str, index_property_name: str
-    ) -> None:
+    async def create_vector_index(self, index_name: str, index_property_name: str) -> None:
         """Create a vector index for a specific property.
 
         Args:
@@ -301,17 +285,15 @@ class RedisAdapter(VectorDBInterface):
             [
                 RedisDataPoint(
                     id=data_point.id,
-                    text=getattr(
-                        data_point, data_point.metadata.get("index_fields", ["text"])[0]
-                    ),
+                    text=getattr(data_point, data_point.metadata.get("index_fields", ["text"])[0]),
                 )
                 for data_point in data_points
             ],
         )
 
     async def retrieve(
-        self, collection_name: str, data_point_ids: List[str]
-    ) -> List[Dict[str, Any]]:
+        self, collection_name: str, data_point_ids: list[str]
+    ) -> list[dict[str, Any]]:
         """Retrieve data points by their IDs.
 
         Args:
@@ -342,17 +324,16 @@ class RedisAdapter(VectorDBInterface):
         except Exception as e:
             logger.error(f"Error retrieving data points: {str(e)}")
             return []
-        finally:
-            await index.disconnect()
 
     async def search(
         self,
         collection_name: str,
-        query_text: Optional[str] = None,
-        query_vector: Optional[List[float]] = None,
-        limit: int = 15,
+        query_text: str | None = None,
+        query_vector: list[float] | None = None,
+        limit: int | None = 15,
         with_vector: bool = False,
-    ) -> List[ScoredResult]:
+        include_payload: bool = True,
+    ) -> list[ScoredResult]:
         """Search for similar vectors in the collection.
 
         Args:
@@ -361,6 +342,7 @@ class RedisAdapter(VectorDBInterface):
             query_vector: Pre-computed query vector.
             limit: Maximum number of results to return.
             with_vector: Whether to include vectors in results.
+            include_payload: Whether to include payloads in results.
 
         Returns:
             List of ScoredResult objects sorted by similarity.
@@ -380,11 +362,11 @@ class RedisAdapter(VectorDBInterface):
 
         index = self._get_index(collection_name)
 
-        if limit == 0:
+        if limit is None:
             info = await index.info()
             limit = info["num_docs"]
 
-        if limit == 0:
+        if limit <= 0:
             return []
 
         try:
@@ -402,7 +384,10 @@ class RedisAdapter(VectorDBInterface):
             )
 
             # Set return fields
-            return_fields = ["id", "text", "payload_data"]
+            if include_payload:
+                return_fields = ["id", "text", "payload_data"]
+            else:
+                return_fields = ["id", "text"]
             if with_vector:
                 return_fields.append("vector")
             vector_query = vector_query.return_fields(*return_fields)
@@ -424,9 +409,7 @@ class RedisAdapter(VectorDBInterface):
                     ScoredResult(
                         id=parse_id(doc["id"].split(":", 1)[1]),
                         payload=payload,
-                        score=float(
-                            doc.get("vector_distance", 0.0)
-                        ),  # RedisVL returns distance
+                        score=float(doc.get("vector_distance", 0.0)),  # RedisVL returns distance
                     )
                 )
             return scored_results
@@ -434,16 +417,15 @@ class RedisAdapter(VectorDBInterface):
         except Exception as e:
             logger.error(f"Error during search: {str(e)}")
             raise e
-        finally:
-            await index.disconnect()
 
     async def batch_search(
         self,
         collection_name: str,
-        query_texts: List[str],
-        limit: Optional[int] = None,
+        query_texts: list[str],
+        limit: int | None,
         with_vectors: bool = False,
-    ) -> List[List[ScoredResult]]:
+        include_payload: bool = True,
+    ) -> list[list[ScoredResult]]:
         """Perform batch search for multiple queries.
 
         Args:
@@ -451,6 +433,7 @@ class RedisAdapter(VectorDBInterface):
             query_texts: List of text queries to search for.
             limit: Maximum number of results per query.
             with_vectors: Whether to include vectors in results.
+            include_payload: Whether to include payloads in results.
 
         Returns:
             List of search results for each query, filtered by score threshold.
@@ -466,6 +449,7 @@ class RedisAdapter(VectorDBInterface):
                 query_vector=vector,
                 limit=limit,
                 with_vector=with_vectors,
+                include_payload=include_payload,
             )
             for vector in vectors
         ]
@@ -474,13 +458,12 @@ class RedisAdapter(VectorDBInterface):
 
         # Filter results by score threshold (Redis uses distance, so lower is better)
         return [
-            [result for result in result_group if result.score < 0.1]
-            for result_group in results
+            [result for result in result_group if result.score < 0.1] for result_group in results
         ]
 
     async def delete_data_points(
-        self, collection_name: str, data_point_ids: List[str]
-    ) -> Dict[str, int]:
+        self, collection_name: str, data_point_ids: list[str]
+    ) -> dict[str, int]:
         """Delete data points by their IDs.
 
         Args:
@@ -496,15 +479,11 @@ class RedisAdapter(VectorDBInterface):
         index = self._get_index(collection_name)
         try:
             deleted_count = await index.drop_documents(data_point_ids)
-            logger.info(
-                f"Deleted {deleted_count} data points from collection {collection_name}"
-            )
+            logger.info(f"Deleted {deleted_count} data points from collection {collection_name}")
             return {"deleted": deleted_count}
         except Exception as e:
             logger.error(f"Error deleting data points: {str(e)}")
             raise e
-        finally:
-            await index.disconnect()
 
     async def prune(self) -> None:
         """Remove all collections and data from Redis.
@@ -521,10 +500,8 @@ class RedisAdapter(VectorDBInterface):
                     if await index.exists():
                         await index.delete(drop=True)
                         logger.info(f"Dropped index {collection_name}")
-                        await index.disconnect()
                 except Exception as e:
                     logger.warning(f"Failed to drop index {collection_name}: {str(e)}")
-                    await index.disconnect()
 
             # Clear the indices cache
             self._indices.clear()
@@ -534,3 +511,12 @@ class RedisAdapter(VectorDBInterface):
         except Exception as e:
             logger.error(f"Error during prune: {str(e)}")
             raise e
+
+    async def get_collection_names(self):
+        """
+        Get names of all collections in the database.
+
+        Returns:
+            List of collection names. In this case of Redis, the return type is a dict.
+        """
+        return self._indices.keys()

{cognee_community_vector_adapter_redis-0.0.3 → cognee_community_vector_adapter_redis-0.1.0}/examples/example.py

@@ -1,5 +1,5 @@
-import os
 import asyncio
+import os
 import pathlib
 from os import path
 
@@ -8,10 +8,11 @@ os.environ.setdefault("LLM_API_KEY", "your-api-key")
 
 
 async def main():
-    from cognee import config, prune, add, cognify, search, SearchType
+    from cognee import SearchType, add, cognify, config, prune, search
 
     # NOTE: Importing the register module we let cognee know it can use the Redis vector adapter
-    from cognee_community_vector_adapter_redis import register
+    # NOTE: The "noqa: F401" mark is to make sure the linter doesn't flag this as an unused import
+    from cognee_community_vector_adapter_redis import register  # noqa: F401
 
     system_path = pathlib.Path(__file__).parent
     config.system_root_directory(path.join(system_path, ".cognee-system"))
@@ -36,16 +37,14 @@ async def main():
 
     await add("""
     Sandwhiches are best served toasted with cheese, ham, mayo,
-    lettuce, mustard, and salt & pepper.          
+    lettuce, mustard, and salt & pepper.
     """)
 
     await cognify()
 
     query_text = "Tell me about NLP"
 
-    search_results = await search(
-        query_type=SearchType.GRAPH_COMPLETION, query_text=query_text
-    )
+    search_results = await search(query_type=SearchType.GRAPH_COMPLETION, query_text=query_text)
 
     for result_text in search_results:
         print("\nSearch result: \n" + result_text)