gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Rework limit=0 for vector adapters (#1450)

Browse source 
<!-- .github/pull_request_template.md -->

## Description
<!-- 
Please provide a clear, human-generated description of the changes in
this PR.
DO NOT use AI-generated descriptions. We want to understand your thought
process and reasoning.
-->
Until now, limit=0 in vector search meant that there is no limit and we
should return everything. This caused confusion and errors, so now it is
reworked so that limit=None means no limit on the search. If someone
puts limit=0, there will be no results returned, as it makes more sense
and is less error prone.

## Type of Change
<!-- Please check the relevant option -->
- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Code refactoring
- [ ] Performance improvement
- [ ] Other (please specify):

## Changes Made
<!-- List the specific changes made in this PR -->
- 
- 
- 

## Testing
<!-- Describe how you tested your changes -->

## Screenshots/Videos (if applicable)
<!-- Add screenshots or videos to help explain your changes -->

## Pre-submission Checklist
<!-- Please check all boxes that apply before submitting your PR -->
- [x] **I have tested my changes thoroughly before submitting this PR**
- [x] **This PR contains minimal changes necessary to address the
issue/feature**
- [x] My code follows the project's coding standards and style
guidelines
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have added necessary documentation (if applicable)
- [ ] All new and existing tests pass
- [ ] I have searched existing PRs to ensure this change hasn't been
submitted already
- [ ] I have linked any relevant issues in the description
- [ ] My commits have clear and descriptive messages

## Related Issues
<!-- Link any related issues using "Fixes #issue_number" or "Relates to
#issue_number" -->

## Additional Notes
<!-- Add any additional notes, concerns, or context for reviewers -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
This commit is contained in:
Vasilije 2025-09-25 21:13:41 +02:00 • committed by GitHub
commit 235f28aefe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 344 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
.github/workflows/vector_db_tests.yml vendored
View file

@ -101,3 +101,30 @@ jobs:
EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
run: uv run python ./cognee/tests/test_pgvector.py
run-lancedb-tests:
name: LanceDB Tests
runs-on: ubuntu-22.04
steps:
- name: Check out
uses: actions/checkout@v4
with:
fetch-depth: 0
- name: Cognee Setup
uses: ./.github/actions/cognee_setup
with:
python-version: ${{ inputs.python-version }}
- name: Run LanceDB Tests
env:
ENV: 'dev'
LLM_MODEL: ${{ secrets.LLM_MODEL }}
LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
run: uv run python ./cognee/tests/test_lancedb.py

6
cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py
View file

@ -234,7 +234,7 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
collection_name: str,
query_text: Optional[str] = None,
query_vector: Optional[List[float]] = None,
limit: int = None,
limit: Optional[int] = None,
with_vector: bool = False,
):
"""
@ -265,10 +265,10 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface):
"Use this option only when vector data is required."
)
# In the case of excessive limit, or zero / negative value, limit will be set to 10.
# In the case of excessive limit, or None / zero / negative value, limit will be set to 10.
if not limit or limit <= self._TOPK_LOWER_BOUND or limit > self._TOPK_UPPER_BOUND:
logger.warning(
"Provided limit (%s) is invalid (zero, negative, or exceeds maximum). "
"Provided limit (%s) is invalid (None, zero, negative, or exceeds maximum). "
"Defaulting to limit=10.",
limit,
)

10
cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py
View file

@ -352,7 +352,7 @@ class ChromaDBAdapter(VectorDBInterface):
collection_name: str,
query_text: str = None,
query_vector: List[float] = None,
limit: int = 15,
limit: Optional[int] = 15,
with_vector: bool = False,
normalized: bool = True,
):
@ -386,9 +386,13 @@ class ChromaDBAdapter(VectorDBInterface):
try:
collection = await self.get_collection(collection_name)
if limit == 0:
if limit is None:
limit = await collection.count()
# If limit is still 0, no need to do the search, just return empty results
if limit <= 0:
return []
results = await collection.query(
query_embeddings=[query_vector],
include=["metadatas", "distances", "embeddings"]
@ -428,7 +432,7 @@ class ChromaDBAdapter(VectorDBInterface):
for row in vector_list
]
except Exception as e:
logger.error(f"Error in search: {str(e)}")
logger.warning(f"Error in search: {str(e)}")
return []
async def batch_search(

8
cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py
View file

@ -223,7 +223,7 @@ class LanceDBAdapter(VectorDBInterface):
collection_name: str,
query_text: str = None,
query_vector: List[float] = None,
limit: int = 15,
limit: Optional[int] = 15,
with_vector: bool = False,
normalized: bool = True,
):
@ -235,11 +235,11 @@ class LanceDBAdapter(VectorDBInterface):
collection = await self.get_collection(collection_name)
if limit == 0:
if limit is None:
limit = await collection.count_rows()
# LanceDB search will break if limit is 0 so we must return
if limit == 0:
if limit <= 0:
return []
results = await collection.vector_search(query_vector).limit(limit).to_pandas()
@ -264,7 +264,7 @@ class LanceDBAdapter(VectorDBInterface):
self,
collection_name: str,
query_texts: List[str],
limit: int = None,
limit: Optional[int] = None,
with_vectors: bool = False,
):
query_vectors = await self.embedding_engine.embed_text(query_texts)

14
cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
View file

@ -3,7 +3,7 @@ from typing import List, Optional, get_type_hints
from sqlalchemy.inspection import inspect
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy import JSON, Column, Table, select, delete, MetaData
from sqlalchemy import JSON, Column, Table, select, delete, MetaData, func
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
from sqlalchemy.exc import ProgrammingError
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
@ -299,7 +299,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
collection_name: str,
query_text: Optional[str] = None,
query_vector: Optional[List[float]] = None,
limit: int = 15,
limit: Optional[int] = 15,
with_vector: bool = False,
) -> List[ScoredResult]:
if query_text is None and query_vector is None:
@ -311,6 +311,16 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
# Get PGVectorDataPoint Table from database
PGVectorDataPoint = await self.get_table(collection_name)
if limit is None:
async with self.get_async_session() as session:
query = select(func.count()).select_from(PGVectorDataPoint)
result = await session.execute(query)
limit = result.scalar_one()
# If limit is still 0, no need to do the search, just return empty results
if limit <= 0:
return []
# NOTE: This needs to be initialized in case search doesn't return a value
closest_items = []

12
cognee/infrastructure/databases/vector/vector_db_interface.py
View file

@ -83,7 +83,7 @@ class VectorDBInterface(Protocol):
collection_name: str,
query_text: Optional[str],
query_vector: Optional[List[float]],
limit: int,
limit: Optional[int],
with_vector: bool = False,
):
"""
@ -98,7 +98,7 @@ class VectorDBInterface(Protocol):
collection.
- query_vector (Optional[List[float]]): An optional vector representation for
searching the collection.
- limit (int): The maximum number of results to return from the search.
- limit (Optional[int]): The maximum number of results to return from the search.
- with_vector (bool): Whether to return the vector representations with search
results. (default False)
"""
@ -106,7 +106,11 @@ class VectorDBInterface(Protocol):
@abstractmethod
async def batch_search(
self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False
self,
collection_name: str,
query_texts: List[str],
limit: Optional[int],
with_vectors: bool = False,
):
"""
Perform a batch search using multiple text queries against a collection.
@ -116,7 +120,7 @@ class VectorDBInterface(Protocol):
- collection_name (str): The name of the collection to conduct the batch search in.
- query_texts (List[str]): A list of text queries to use for the search.
- limit (int): The maximum number of results to return for each query.
- limit (Optional[int]): The maximum number of results to return for each query.
- with_vectors (bool): Whether to include vector representations with search
results. (default False)
"""

2
cognee/modules/graph/cognee_graph/CogneeGraph.py
View file

@ -161,7 +161,7 @@ class CogneeGraph(CogneeAbstractGraph):
edge_distances = await vector_engine.search(
collection_name="EdgeType_relationship_name",
query_vector=query_vector,
limit=0,
limit=None,
)
projection_time = time.time() - start_time
logger.info(

2
cognee/modules/retrieval/insights_retriever.py
View file

@ -25,7 +25,7 @@ class InsightsRetriever(BaseGraphRetriever):
- top_k
"""
def __init__(self, exploration_levels: int = 1, top_k: int = 5):
def __init__(self, exploration_levels: int = 1, top_k: Optional[int] = 5):
"""Initialize retriever with exploration levels and search parameters."""
self.exploration_levels = exploration_levels
self.top_k = top_k

2
cognee/modules/retrieval/temporal_retriever.py
View file

@ -129,7 +129,7 @@ class TemporalRetriever(GraphCompletionRetriever):
query_vector = (await vector_engine.embedding_engine.embed_text([query]))[0]
vector_search_results = await vector_engine.search(
collection_name="Event_name", query_vector=query_vector, limit=0
collection_name="Event_name", query_vector=query_vector, limit=None
)
top_k_events = await self.filter_top_k_events(relevant_events, vector_search_results)

2
cognee/modules/retrieval/utils/brute_force_triplet_search.py
View file

@ -144,7 +144,7 @@ async def brute_force_triplet_search(
async def search_in_collection(collection_name: str):
try:
return await vector_engine.search(
collection_name=collection_name, query_vector=query_vector, limit=0
collection_name=collection_name, query_vector=query_vector, limit=None
)
except CollectionNotFoundError:
return []

35
cognee/tests/test_chromadb.py
View file

@ -67,6 +67,39 @@ async def test_getting_of_documents(dataset_name_1):
)
async def test_vector_engine_search_none_limit():
file_path = os.path.join(
pathlib.Path(__file__).resolve().parent.parent.parent,
"examples",
"data",
"alice_in_wonderland.txt",
)
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
await cognee.add(file_path)
await cognee.cognify()
query_text = "List me all the important characters in Alice in Wonderland."
from cognee.infrastructure.databases.vector import get_vector_engine
vector_engine = get_vector_engine()
collection_name = "Entity_name"
query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0]
result = await vector_engine.search(
collection_name=collection_name, query_vector=query_vector, limit=None
)
# Check that we did not accidentally use any default value for limit in vector search along the way (like 5, 10, or 15)
assert len(result) > 15
async def main():
cognee.config.set_vector_db_config(
{
@ -165,6 +198,8 @@ async def main():
tables_in_database = await vector_engine.get_collection_names()
assert len(tables_in_database) == 0, "ChromaDB database is not empty"
await test_vector_engine_search_none_limit()
if __name__ == "__main__":
import asyncio

206
cognee/tests/test_lancedb.py Normal file
View file

@ -0,0 +1,206 @@
import os
import pathlib
import cognee
from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.files.storage import get_storage_config
from cognee.modules.data.models import Data
from cognee.modules.users.methods import get_default_user
from cognee.modules.search.types import SearchType
from cognee.modules.search.operations import get_history
logger = get_logger()
async def test_local_file_deletion(data_text, file_location):
from sqlalchemy import select
import hashlib
from cognee.infrastructure.databases.relational import get_relational_engine
engine = get_relational_engine()
async with engine.get_async_session() as session:
# Get hash of data contents
encoded_text = data_text.encode("utf-8")
data_hash = hashlib.md5(encoded_text).hexdigest()
# Get data entry from database based on hash contents
data = (await session.scalars(select(Data).where(Data.content_hash == data_hash))).one()
assert os.path.isfile(data.raw_data_location.replace("file://", "")), (
f"Data location doesn't exist: {data.raw_data_location}"
)
# Test deletion of data along with local files created by cognee
await engine.delete_data_entity(data.id)
assert not os.path.exists(data.raw_data_location.replace("file://", "")), (
f"Data location still exists after deletion: {data.raw_data_location}"
)
async with engine.get_async_session() as session:
# Get data entry from database based on file path
data = (
await session.scalars(select(Data).where(Data.raw_data_location == file_location))
).one()
assert os.path.isfile(data.raw_data_location.replace("file://", "")), (
f"Data location doesn't exist: {data.raw_data_location}"
)
# Test local files not created by cognee won't get deleted
await engine.delete_data_entity(data.id)
assert os.path.exists(data.raw_data_location.replace("file://", "")), (
f"Data location doesn't exists: {data.raw_data_location}"
)
async def test_getting_of_documents(dataset_name_1):
# Test getting of documents for search per dataset
from cognee.modules.users.permissions.methods import get_document_ids_for_user
user = await get_default_user()
document_ids = await get_document_ids_for_user(user.id, [dataset_name_1])
assert len(document_ids) == 1, (
f"Number of expected documents doesn't match {len(document_ids)} != 1"
)
# Test getting of documents for search when no dataset is provided
user = await get_default_user()
document_ids = await get_document_ids_for_user(user.id)
assert len(document_ids) == 2, (
f"Number of expected documents doesn't match {len(document_ids)} != 2"
)
async def test_vector_engine_search_none_limit():
file_path = os.path.join(
pathlib.Path(__file__).resolve().parent.parent.parent,
"examples",
"data",
"alice_in_wonderland.txt",
)
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
await cognee.add(file_path)
await cognee.cognify()
query_text = "List me all the important characters in Alice in Wonderland."
from cognee.infrastructure.databases.vector import get_vector_engine
vector_engine = get_vector_engine()
collection_name = "Entity_name"
query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0]
result = await vector_engine.search(
collection_name=collection_name, query_vector=query_vector, limit=None
)
# Check that we did not accidentally use any default value for limit in vector search along the way (like 5, 10, or 15)
assert len(result) > 15
async def main():
cognee.config.set_vector_db_config(
{
"vector_db_provider": "lancedb",
}
)
data_directory_path = str(
pathlib.Path(
os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_lancedb")
).resolve()
)
cognee.config.data_root_directory(data_directory_path)
cognee_directory_path = str(
pathlib.Path(
os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_lancedb")
).resolve()
)
cognee.config.system_root_directory(cognee_directory_path)
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
dataset_name_1 = "natural_language"
dataset_name_2 = "quantum"
explanation_file_path = os.path.join(
pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
)
await cognee.add([explanation_file_path], dataset_name_1)
text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
"""
await cognee.add([text], dataset_name_2)
await cognee.cognify([dataset_name_2, dataset_name_1])
from cognee.infrastructure.databases.vector import get_vector_engine
await test_getting_of_documents(dataset_name_1)
vector_engine = get_vector_engine()
random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0]
random_node_name = random_node.payload["text"]
search_results = await cognee.search(
query_type=SearchType.INSIGHTS, query_text=random_node_name
)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted sentences are:\n")
for result in search_results:
print(f"{result}\n")
search_results = await cognee.search(
query_type=SearchType.CHUNKS, query_text=random_node_name, datasets=[dataset_name_2]
)
assert len(search_results) != 0, "The search results list is empty."
print("\n\nExtracted chunks are:\n")
for result in search_results:
print(f"{result}\n")
graph_completion = await cognee.search(
query_type=SearchType.GRAPH_COMPLETION,
query_text=random_node_name,
datasets=[dataset_name_2],
)
assert len(graph_completion) != 0, "Completion result is empty."
print("Completion result is:")
print(graph_completion)
search_results = await cognee.search(
query_type=SearchType.SUMMARIES, query_text=random_node_name
)
assert len(search_results) != 0, "Query related summaries don't exist."
print("\n\nExtracted summaries are:\n")
for result in search_results:
print(f"{result}\n")
user = await get_default_user()
history = await get_history(user.id)
assert len(history) == 8, "Search history is not correct."
await cognee.prune.prune_data()
data_root_directory = get_storage_config()["data_root_directory"]
assert not os.path.isdir(data_root_directory), "Local data files are not deleted"
await cognee.prune.prune_system(metadata=True)
connection = await vector_engine.get_connection()
tables_in_database = await connection.table_names()
assert len(tables_in_database) == 0, "LanceDB database is not empty"
await test_vector_engine_search_none_limit()
if __name__ == "__main__":
import asyncio
asyncio.run(main())

35
cognee/tests/test_pgvector.py
View file

@ -68,6 +68,39 @@ async def test_getting_of_documents(dataset_name_1):
)
async def test_vector_engine_search_none_limit():
file_path = os.path.join(
pathlib.Path(__file__).resolve().parent.parent.parent,
"examples",
"data",
"alice_in_wonderland.txt",
)
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
await cognee.add(file_path)
await cognee.cognify()
query_text = "List me all the important characters in Alice in Wonderland."
from cognee.infrastructure.databases.vector import get_vector_engine
vector_engine = get_vector_engine()
collection_name = "Entity_name"
query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0]
result = await vector_engine.search(
collection_name=collection_name, query_vector=query_vector, limit=None
)
# Check that we did not accidentally use any default value for limit in vector search along the way (like 5, 10, or 15)
assert len(result) > 15
async def main():
cognee.config.set_vector_db_config(
{"vector_db_url": "", "vector_db_key": "", "vector_db_provider": "pgvector"}
@ -174,6 +207,8 @@ async def main():
tables_in_database = await vector_engine.get_table_names()
assert len(tables_in_database) == 0, "PostgreSQL database is not empty"
await test_vector_engine_search_none_limit()
if __name__ == "__main__":
import asyncio

3
examples/python/cognee_simple_document_demo.py
View file

@ -15,6 +15,9 @@ async def cognee_demo():
current_directory = Path(__file__).resolve().parent.parent
file_path = os.path.join(current_directory, "data", "alice_in_wonderland.txt")
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
# Call Cognee to process document
await cognee.add(file_path)
await cognee.cognify()