cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- cognee/__init__.py +1 -0
- cognee/__main__.py +4 -0
- cognee/api/client.py +28 -3
- cognee/api/health.py +10 -13
- cognee/api/v1/add/add.py +20 -6
- cognee/api/v1/add/routers/get_add_router.py +12 -37
- cognee/api/v1/cloud/routers/__init__.py +1 -0
- cognee/api/v1/cloud/routers/get_checks_router.py +23 -0
- cognee/api/v1/cognify/code_graph_pipeline.py +14 -3
- cognee/api/v1/cognify/cognify.py +67 -105
- cognee/api/v1/cognify/routers/get_cognify_router.py +11 -3
- cognee/api/v1/datasets/routers/get_datasets_router.py +16 -5
- cognee/api/v1/memify/routers/__init__.py +1 -0
- cognee/api/v1/memify/routers/get_memify_router.py +100 -0
- cognee/api/v1/notebooks/routers/__init__.py +1 -0
- cognee/api/v1/notebooks/routers/get_notebooks_router.py +96 -0
- cognee/api/v1/responses/default_tools.py +4 -0
- cognee/api/v1/responses/dispatch_function.py +6 -1
- cognee/api/v1/responses/models.py +1 -1
- cognee/api/v1/search/routers/get_search_router.py +20 -1
- cognee/api/v1/search/search.py +17 -4
- cognee/api/v1/sync/__init__.py +17 -0
- cognee/api/v1/sync/routers/__init__.py +3 -0
- cognee/api/v1/sync/routers/get_sync_router.py +241 -0
- cognee/api/v1/sync/sync.py +877 -0
- cognee/api/v1/users/routers/get_auth_router.py +13 -1
- cognee/base_config.py +10 -1
- cognee/cli/__init__.py +10 -0
- cognee/cli/_cognee.py +180 -0
- cognee/cli/commands/__init__.py +1 -0
- cognee/cli/commands/add_command.py +80 -0
- cognee/cli/commands/cognify_command.py +128 -0
- cognee/cli/commands/config_command.py +225 -0
- cognee/cli/commands/delete_command.py +80 -0
- cognee/cli/commands/search_command.py +149 -0
- cognee/cli/config.py +33 -0
- cognee/cli/debug.py +21 -0
- cognee/cli/echo.py +45 -0
- cognee/cli/exceptions.py +23 -0
- cognee/cli/minimal_cli.py +97 -0
- cognee/cli/reference.py +26 -0
- cognee/cli/suppress_logging.py +12 -0
- cognee/eval_framework/corpus_builder/corpus_builder_executor.py +2 -2
- cognee/eval_framework/eval_config.py +1 -1
- cognee/infrastructure/databases/graph/config.py +10 -4
- cognee/infrastructure/databases/graph/get_graph_engine.py +4 -9
- cognee/infrastructure/databases/graph/kuzu/adapter.py +199 -2
- cognee/infrastructure/databases/graph/neo4j_driver/adapter.py +138 -0
- cognee/infrastructure/databases/relational/__init__.py +2 -0
- cognee/infrastructure/databases/relational/get_async_session.py +15 -0
- cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py +6 -1
- cognee/infrastructure/databases/relational/with_async_session.py +25 -0
- cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py +1 -1
- cognee/infrastructure/databases/vector/config.py +13 -6
- cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py +6 -4
- cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py +16 -7
- cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py +5 -5
- cognee/infrastructure/databases/vector/embeddings/config.py +2 -2
- cognee/infrastructure/databases/vector/embeddings/embedding_rate_limiter.py +2 -6
- cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py +10 -7
- cognee/infrastructure/files/storage/LocalFileStorage.py +9 -0
- cognee/infrastructure/files/storage/S3FileStorage.py +5 -0
- cognee/infrastructure/files/storage/StorageManager.py +7 -1
- cognee/infrastructure/files/storage/storage.py +16 -0
- cognee/infrastructure/files/utils/get_data_file_path.py +14 -9
- cognee/infrastructure/files/utils/get_file_metadata.py +2 -1
- cognee/infrastructure/llm/LLMGateway.py +32 -5
- cognee/infrastructure/llm/config.py +6 -4
- cognee/infrastructure/llm/prompts/extract_query_time.txt +15 -0
- cognee/infrastructure/llm/prompts/generate_event_entity_prompt.txt +25 -0
- cognee/infrastructure/llm/prompts/generate_event_graph_prompt.txt +30 -0
- cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/knowledge_graph/extract_content_graph.py +16 -5
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/__init__.py +2 -0
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_event_entities.py +44 -0
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/__init__.py +1 -0
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_content_graph.py +19 -15
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_event_graph.py +46 -0
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py +3 -3
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py +3 -3
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py +2 -2
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py +14 -8
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py +6 -4
- cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py +28 -4
- cognee/infrastructure/llm/tokenizer/Gemini/adapter.py +2 -2
- cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py +3 -3
- cognee/infrastructure/llm/tokenizer/Mistral/adapter.py +3 -3
- cognee/infrastructure/llm/tokenizer/TikToken/adapter.py +6 -6
- cognee/infrastructure/llm/utils.py +7 -7
- cognee/infrastructure/utils/run_sync.py +8 -1
- cognee/modules/chunking/models/DocumentChunk.py +4 -3
- cognee/modules/cloud/exceptions/CloudApiKeyMissingError.py +15 -0
- cognee/modules/cloud/exceptions/CloudConnectionError.py +15 -0
- cognee/modules/cloud/exceptions/__init__.py +2 -0
- cognee/modules/cloud/operations/__init__.py +1 -0
- cognee/modules/cloud/operations/check_api_key.py +25 -0
- cognee/modules/data/deletion/prune_system.py +1 -1
- cognee/modules/data/methods/__init__.py +2 -0
- cognee/modules/data/methods/check_dataset_name.py +1 -1
- cognee/modules/data/methods/create_authorized_dataset.py +19 -0
- cognee/modules/data/methods/get_authorized_dataset.py +11 -5
- cognee/modules/data/methods/get_authorized_dataset_by_name.py +16 -0
- cognee/modules/data/methods/get_dataset_data.py +1 -1
- cognee/modules/data/methods/load_or_create_datasets.py +2 -20
- cognee/modules/engine/models/Event.py +16 -0
- cognee/modules/engine/models/Interval.py +8 -0
- cognee/modules/engine/models/Timestamp.py +13 -0
- cognee/modules/engine/models/__init__.py +3 -0
- cognee/modules/engine/utils/__init__.py +2 -0
- cognee/modules/engine/utils/generate_event_datapoint.py +46 -0
- cognee/modules/engine/utils/generate_timestamp_datapoint.py +51 -0
- cognee/modules/graph/cognee_graph/CogneeGraph.py +2 -2
- cognee/modules/graph/methods/get_formatted_graph_data.py +3 -2
- cognee/modules/graph/utils/__init__.py +1 -0
- cognee/modules/graph/utils/resolve_edges_to_text.py +71 -0
- cognee/modules/memify/__init__.py +1 -0
- cognee/modules/memify/memify.py +118 -0
- cognee/modules/notebooks/methods/__init__.py +5 -0
- cognee/modules/notebooks/methods/create_notebook.py +26 -0
- cognee/modules/notebooks/methods/delete_notebook.py +13 -0
- cognee/modules/notebooks/methods/get_notebook.py +21 -0
- cognee/modules/notebooks/methods/get_notebooks.py +18 -0
- cognee/modules/notebooks/methods/update_notebook.py +17 -0
- cognee/modules/notebooks/models/Notebook.py +53 -0
- cognee/modules/notebooks/models/__init__.py +1 -0
- cognee/modules/notebooks/operations/__init__.py +1 -0
- cognee/modules/notebooks/operations/run_in_local_sandbox.py +55 -0
- cognee/modules/pipelines/__init__.py +1 -1
- cognee/modules/pipelines/exceptions/tasks.py +18 -0
- cognee/modules/pipelines/layers/__init__.py +1 -0
- cognee/modules/pipelines/layers/check_pipeline_run_qualification.py +59 -0
- cognee/modules/pipelines/layers/pipeline_execution_mode.py +127 -0
- cognee/modules/pipelines/layers/reset_dataset_pipeline_run_status.py +28 -0
- cognee/modules/pipelines/layers/resolve_authorized_user_dataset.py +34 -0
- cognee/modules/pipelines/layers/resolve_authorized_user_datasets.py +55 -0
- cognee/modules/pipelines/layers/setup_and_check_environment.py +41 -0
- cognee/modules/pipelines/layers/validate_pipeline_tasks.py +20 -0
- cognee/modules/pipelines/methods/__init__.py +2 -0
- cognee/modules/pipelines/methods/get_pipeline_runs_by_dataset.py +34 -0
- cognee/modules/pipelines/methods/reset_pipeline_run_status.py +16 -0
- cognee/modules/pipelines/operations/__init__.py +0 -1
- cognee/modules/pipelines/operations/log_pipeline_run_initiated.py +1 -1
- cognee/modules/pipelines/operations/pipeline.py +24 -138
- cognee/modules/pipelines/operations/run_tasks.py +17 -41
- cognee/modules/retrieval/base_feedback.py +11 -0
- cognee/modules/retrieval/base_graph_retriever.py +18 -0
- cognee/modules/retrieval/base_retriever.py +1 -1
- cognee/modules/retrieval/code_retriever.py +8 -0
- cognee/modules/retrieval/coding_rules_retriever.py +31 -0
- cognee/modules/retrieval/completion_retriever.py +9 -3
- cognee/modules/retrieval/context_providers/TripletSearchContextProvider.py +1 -0
- cognee/modules/retrieval/cypher_search_retriever.py +1 -9
- cognee/modules/retrieval/graph_completion_context_extension_retriever.py +29 -13
- cognee/modules/retrieval/graph_completion_cot_retriever.py +30 -13
- cognee/modules/retrieval/graph_completion_retriever.py +107 -56
- cognee/modules/retrieval/graph_summary_completion_retriever.py +5 -1
- cognee/modules/retrieval/insights_retriever.py +14 -3
- cognee/modules/retrieval/natural_language_retriever.py +0 -4
- cognee/modules/retrieval/summaries_retriever.py +1 -1
- cognee/modules/retrieval/temporal_retriever.py +152 -0
- cognee/modules/retrieval/user_qa_feedback.py +83 -0
- cognee/modules/retrieval/utils/brute_force_triplet_search.py +7 -32
- cognee/modules/retrieval/utils/completion.py +10 -3
- cognee/modules/retrieval/utils/extract_uuid_from_node.py +18 -0
- cognee/modules/retrieval/utils/models.py +40 -0
- cognee/modules/search/methods/get_search_type_tools.py +168 -0
- cognee/modules/search/methods/no_access_control_search.py +47 -0
- cognee/modules/search/methods/search.py +239 -118
- cognee/modules/search/types/SearchResult.py +21 -0
- cognee/modules/search/types/SearchType.py +3 -0
- cognee/modules/search/types/__init__.py +1 -0
- cognee/modules/search/utils/__init__.py +2 -0
- cognee/modules/search/utils/prepare_search_result.py +41 -0
- cognee/modules/search/utils/transform_context_to_graph.py +38 -0
- cognee/modules/settings/get_settings.py +2 -2
- cognee/modules/sync/__init__.py +1 -0
- cognee/modules/sync/methods/__init__.py +23 -0
- cognee/modules/sync/methods/create_sync_operation.py +53 -0
- cognee/modules/sync/methods/get_sync_operation.py +107 -0
- cognee/modules/sync/methods/update_sync_operation.py +248 -0
- cognee/modules/sync/models/SyncOperation.py +142 -0
- cognee/modules/sync/models/__init__.py +3 -0
- cognee/modules/users/__init__.py +0 -1
- cognee/modules/users/methods/__init__.py +4 -1
- cognee/modules/users/methods/create_user.py +26 -1
- cognee/modules/users/methods/get_authenticated_user.py +36 -42
- cognee/modules/users/methods/get_default_user.py +3 -1
- cognee/modules/users/permissions/methods/get_specific_user_permission_datasets.py +2 -1
- cognee/root_dir.py +19 -0
- cognee/shared/CodeGraphEntities.py +1 -0
- cognee/shared/logging_utils.py +143 -32
- cognee/shared/utils.py +0 -1
- cognee/tasks/codingagents/coding_rule_associations.py +127 -0
- cognee/tasks/graph/extract_graph_from_data.py +6 -2
- cognee/tasks/ingestion/save_data_item_to_storage.py +23 -0
- cognee/tasks/memify/__init__.py +2 -0
- cognee/tasks/memify/extract_subgraph.py +7 -0
- cognee/tasks/memify/extract_subgraph_chunks.py +11 -0
- cognee/tasks/repo_processor/get_local_dependencies.py +2 -0
- cognee/tasks/repo_processor/get_repo_file_dependencies.py +144 -47
- cognee/tasks/storage/add_data_points.py +33 -3
- cognee/tasks/temporal_graph/__init__.py +1 -0
- cognee/tasks/temporal_graph/add_entities_to_event.py +85 -0
- cognee/tasks/temporal_graph/enrich_events.py +34 -0
- cognee/tasks/temporal_graph/extract_events_and_entities.py +32 -0
- cognee/tasks/temporal_graph/extract_knowledge_graph_from_events.py +41 -0
- cognee/tasks/temporal_graph/models.py +49 -0
- cognee/tests/integration/cli/__init__.py +3 -0
- cognee/tests/integration/cli/test_cli_integration.py +331 -0
- cognee/tests/integration/documents/PdfDocument_test.py +2 -2
- cognee/tests/integration/documents/TextDocument_test.py +2 -4
- cognee/tests/integration/documents/UnstructuredDocument_test.py +5 -8
- cognee/tests/{test_deletion.py → test_delete_hard.py} +0 -37
- cognee/tests/test_delete_soft.py +85 -0
- cognee/tests/test_kuzu.py +2 -2
- cognee/tests/test_neo4j.py +2 -2
- cognee/tests/test_permissions.py +3 -3
- cognee/tests/test_relational_db_migration.py +7 -5
- cognee/tests/test_search_db.py +136 -23
- cognee/tests/test_temporal_graph.py +167 -0
- cognee/tests/unit/api/__init__.py +1 -0
- cognee/tests/unit/api/test_conditional_authentication_endpoints.py +246 -0
- cognee/tests/unit/cli/__init__.py +3 -0
- cognee/tests/unit/cli/test_cli_commands.py +483 -0
- cognee/tests/unit/cli/test_cli_edge_cases.py +625 -0
- cognee/tests/unit/cli/test_cli_main.py +173 -0
- cognee/tests/unit/cli/test_cli_runner.py +62 -0
- cognee/tests/unit/cli/test_cli_utils.py +127 -0
- cognee/tests/unit/modules/retrieval/chunks_retriever_test.py +18 -2
- cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py +12 -15
- cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py +10 -15
- cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py +4 -3
- cognee/tests/unit/modules/retrieval/insights_retriever_test.py +4 -2
- cognee/tests/unit/modules/retrieval/rag_completion_retriever_test.py +18 -2
- cognee/tests/unit/modules/retrieval/temporal_retriever_test.py +225 -0
- cognee/tests/unit/modules/users/__init__.py +1 -0
- cognee/tests/unit/modules/users/test_conditional_authentication.py +277 -0
- cognee/tests/unit/processing/utils/utils_test.py +20 -1
- {cognee-0.2.3.dev1.dist-info → cognee-0.3.0.dev0.dist-info}/METADATA +13 -9
- {cognee-0.2.3.dev1.dist-info → cognee-0.3.0.dev0.dist-info}/RECORD +245 -135
- cognee-0.3.0.dev0.dist-info/entry_points.txt +2 -0
- cognee/infrastructure/databases/graph/networkx/adapter.py +0 -1017
- cognee/infrastructure/pipeline/models/Operation.py +0 -60
- cognee/notebooks/github_analysis_step_by_step.ipynb +0 -37
- cognee/tests/tasks/descriptive_metrics/networkx_metrics_test.py +0 -7
- cognee/tests/unit/modules/search/search_methods_test.py +0 -223
- /cognee/{infrastructure/databases/graph/networkx → api/v1/memify}/__init__.py +0 -0
- /cognee/{infrastructure/pipeline/models → tasks/codingagents}/__init__.py +0 -0
- {cognee-0.2.3.dev1.dist-info → cognee-0.3.0.dev0.dist-info}/WHEEL +0 -0
- {cognee-0.2.3.dev1.dist-info → cognee-0.3.0.dev0.dist-info}/licenses/LICENSE +0 -0
- {cognee-0.2.3.dev1.dist-info → cognee-0.3.0.dev0.dist-info}/licenses/NOTICE.md +0 -0
|
@@ -1,60 +0,0 @@
|
|
|
1
|
-
from datetime import datetime, timezone
|
|
2
|
-
from sqlalchemy.orm import Mapped, MappedColumn
|
|
3
|
-
from sqlalchemy import Column, DateTime, ForeignKey, Enum, JSON
|
|
4
|
-
from cognee.infrastructure.databases.relational import Base, UUID
|
|
5
|
-
|
|
6
|
-
|
|
7
|
-
class OperationType(Enum):
|
|
8
|
-
"""
|
|
9
|
-
Define various types of operations for data handling.
|
|
10
|
-
|
|
11
|
-
Public methods:
|
|
12
|
-
- __str__(): Returns a string representation of the operation type.
|
|
13
|
-
|
|
14
|
-
Instance variables:
|
|
15
|
-
- MERGE_DATA: Represents the merge data operation type.
|
|
16
|
-
- APPEND_DATA: Represents the append data operation type.
|
|
17
|
-
"""
|
|
18
|
-
|
|
19
|
-
MERGE_DATA = "MERGE_DATA"
|
|
20
|
-
APPEND_DATA = "APPEND_DATA"
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
class OperationStatus(Enum):
|
|
24
|
-
"""
|
|
25
|
-
Represent the status of an operation with predefined states.
|
|
26
|
-
"""
|
|
27
|
-
|
|
28
|
-
STARTED = "OPERATION_STARTED"
|
|
29
|
-
IN_PROGRESS = "OPERATION_IN_PROGRESS"
|
|
30
|
-
COMPLETE = "OPERATION_COMPLETE"
|
|
31
|
-
ERROR = "OPERATION_ERROR"
|
|
32
|
-
CANCELLED = "OPERATION_CANCELLED"
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
class Operation(Base):
|
|
36
|
-
"""
|
|
37
|
-
Represents an operation in the system, extending the Base class.
|
|
38
|
-
|
|
39
|
-
This class defines the structure of the 'operation' table, including fields for the
|
|
40
|
-
operation's ID, status, type, associated data, metadata, and creation timestamp. The
|
|
41
|
-
public methods available in this class are inherited from the Base class. Instance
|
|
42
|
-
variables include:
|
|
43
|
-
- id: Unique identifier for the operation.
|
|
44
|
-
- status: The current status of the operation.
|
|
45
|
-
- operation_type: The type of operation being represented.
|
|
46
|
-
- data_id: Foreign key referencing the associated data's ID.
|
|
47
|
-
- meta_data: Additional metadata related to the operation.
|
|
48
|
-
- created_at: Timestamp for when the operation was created.
|
|
49
|
-
"""
|
|
50
|
-
|
|
51
|
-
__tablename__ = "operation"
|
|
52
|
-
|
|
53
|
-
id = Column(UUID, primary_key=True)
|
|
54
|
-
status = Column(Enum(OperationStatus))
|
|
55
|
-
operation_type = Column(Enum(OperationType))
|
|
56
|
-
|
|
57
|
-
data_id = Column(UUID, ForeignKey("data.id"))
|
|
58
|
-
meta_data: Mapped[dict] = MappedColumn(type_=JSON)
|
|
59
|
-
|
|
60
|
-
created_at = Column(DateTime, default=datetime.now(timezone.utc))
|
|
@@ -1,37 +0,0 @@
|
|
|
1
|
-
{
|
|
2
|
-
"cells": [
|
|
3
|
-
{
|
|
4
|
-
"cell_type": "code",
|
|
5
|
-
"execution_count": null,
|
|
6
|
-
"id": "initial_id",
|
|
7
|
-
"metadata": {
|
|
8
|
-
"collapsed": true
|
|
9
|
-
},
|
|
10
|
-
"outputs": [],
|
|
11
|
-
"source": [
|
|
12
|
-
""
|
|
13
|
-
]
|
|
14
|
-
}
|
|
15
|
-
],
|
|
16
|
-
"metadata": {
|
|
17
|
-
"kernelspec": {
|
|
18
|
-
"display_name": "Python 3",
|
|
19
|
-
"language": "python",
|
|
20
|
-
"name": "python3"
|
|
21
|
-
},
|
|
22
|
-
"language_info": {
|
|
23
|
-
"codemirror_mode": {
|
|
24
|
-
"name": "ipython",
|
|
25
|
-
"version": 2
|
|
26
|
-
},
|
|
27
|
-
"file_extension": ".py",
|
|
28
|
-
"mimetype": "text/x-python",
|
|
29
|
-
"name": "python",
|
|
30
|
-
"nbconvert_exporter": "python",
|
|
31
|
-
"pygments_lexer": "ipython2",
|
|
32
|
-
"version": "2.7.6"
|
|
33
|
-
}
|
|
34
|
-
},
|
|
35
|
-
"nbformat": 4,
|
|
36
|
-
"nbformat_minor": 5
|
|
37
|
-
}
|
|
@@ -1,7 +0,0 @@
|
|
|
1
|
-
from cognee.tests.tasks.descriptive_metrics.metrics_test_utils import assert_metrics
|
|
2
|
-
import asyncio
|
|
3
|
-
|
|
4
|
-
|
|
5
|
-
if __name__ == "__main__":
|
|
6
|
-
asyncio.run(assert_metrics(provider="networkx", include_optional=False))
|
|
7
|
-
asyncio.run(assert_metrics(provider="networkx", include_optional=True))
|
|
@@ -1,223 +0,0 @@
|
|
|
1
|
-
import json
|
|
2
|
-
import uuid
|
|
3
|
-
from unittest.mock import AsyncMock, MagicMock, patch
|
|
4
|
-
|
|
5
|
-
import pytest
|
|
6
|
-
from pylint.checkers.utils import node_type
|
|
7
|
-
|
|
8
|
-
from cognee.modules.search.exceptions import UnsupportedSearchTypeError
|
|
9
|
-
from cognee.modules.search.methods.search import search, specific_search
|
|
10
|
-
from cognee.modules.search.types import SearchType
|
|
11
|
-
from cognee.modules.users.models import User
|
|
12
|
-
import sys
|
|
13
|
-
|
|
14
|
-
search_module = sys.modules.get("cognee.modules.search.methods.search")
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
@pytest.fixture
|
|
18
|
-
def mock_user():
|
|
19
|
-
user = MagicMock(spec=User)
|
|
20
|
-
user.id = uuid.uuid4()
|
|
21
|
-
return user
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
@pytest.mark.asyncio
|
|
25
|
-
@patch.object(search_module, "log_query")
|
|
26
|
-
@patch.object(search_module, "log_result")
|
|
27
|
-
@patch.object(search_module, "specific_search")
|
|
28
|
-
async def test_search(
|
|
29
|
-
mock_specific_search,
|
|
30
|
-
mock_log_result,
|
|
31
|
-
mock_log_query,
|
|
32
|
-
mock_user,
|
|
33
|
-
):
|
|
34
|
-
# Setup
|
|
35
|
-
query_text = "test query"
|
|
36
|
-
query_type = SearchType.CHUNKS
|
|
37
|
-
datasets = ["dataset1", "dataset2"]
|
|
38
|
-
|
|
39
|
-
# Mock the query logging
|
|
40
|
-
mock_query = MagicMock()
|
|
41
|
-
mock_query.id = uuid.uuid4()
|
|
42
|
-
mock_log_query.return_value = mock_query
|
|
43
|
-
|
|
44
|
-
# Mock document IDs
|
|
45
|
-
doc_id1 = uuid.uuid4()
|
|
46
|
-
doc_id2 = uuid.uuid4()
|
|
47
|
-
|
|
48
|
-
# Mock search results
|
|
49
|
-
search_results = [
|
|
50
|
-
{"document_id": str(doc_id1), "content": "Result 1"},
|
|
51
|
-
{"document_id": str(doc_id2), "content": "Result 2"},
|
|
52
|
-
]
|
|
53
|
-
mock_specific_search.return_value = search_results
|
|
54
|
-
|
|
55
|
-
# Execute
|
|
56
|
-
await search(query_text, query_type, datasets, mock_user)
|
|
57
|
-
|
|
58
|
-
# Verify
|
|
59
|
-
mock_log_query.assert_called_once_with(query_text, query_type.value, mock_user.id)
|
|
60
|
-
mock_specific_search.assert_called_once_with(
|
|
61
|
-
query_type,
|
|
62
|
-
query_text,
|
|
63
|
-
mock_user,
|
|
64
|
-
system_prompt_path="answer_simple_question.txt",
|
|
65
|
-
top_k=10,
|
|
66
|
-
node_type=None,
|
|
67
|
-
node_name=None,
|
|
68
|
-
)
|
|
69
|
-
|
|
70
|
-
# Verify result logging
|
|
71
|
-
mock_log_result.assert_called_once()
|
|
72
|
-
# Check that the first argument is the query ID
|
|
73
|
-
assert mock_log_result.call_args[0][0] == mock_query.id
|
|
74
|
-
# The second argument should be the JSON string of the filtered results
|
|
75
|
-
# We can't directly compare the JSON strings due to potential ordering differences
|
|
76
|
-
# So we parse the JSON and compare the objects
|
|
77
|
-
logged_results = json.loads(mock_log_result.call_args[0][1])
|
|
78
|
-
assert len(logged_results) == 2
|
|
79
|
-
assert logged_results[0]["document_id"] == str(doc_id1)
|
|
80
|
-
assert logged_results[1]["document_id"] == str(doc_id2)
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
@pytest.mark.asyncio
|
|
84
|
-
@patch.object(search_module, "SummariesRetriever")
|
|
85
|
-
@patch.object(search_module, "send_telemetry")
|
|
86
|
-
async def test_specific_search_summaries(mock_send_telemetry, mock_summaries_retriever, mock_user):
|
|
87
|
-
# Setup
|
|
88
|
-
query = "test query"
|
|
89
|
-
query_type = SearchType.SUMMARIES
|
|
90
|
-
|
|
91
|
-
# Mock the retriever
|
|
92
|
-
mock_retriever = MagicMock()
|
|
93
|
-
mock_retriever.get_completion = AsyncMock()
|
|
94
|
-
mock_retriever.get_completion.return_value = [{"content": "Summary result"}]
|
|
95
|
-
mock_summaries_retriever.return_value = mock_retriever
|
|
96
|
-
|
|
97
|
-
# Execute
|
|
98
|
-
results = await specific_search(query_type, query, mock_user)
|
|
99
|
-
|
|
100
|
-
# Verify
|
|
101
|
-
mock_summaries_retriever.assert_called_once()
|
|
102
|
-
mock_retriever.get_completion.assert_called_once_with(query)
|
|
103
|
-
mock_send_telemetry.assert_called()
|
|
104
|
-
assert len(results) == 1
|
|
105
|
-
assert results[0]["content"] == "Summary result"
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
@pytest.mark.asyncio
|
|
109
|
-
@patch.object(search_module, "InsightsRetriever")
|
|
110
|
-
@patch.object(search_module, "send_telemetry")
|
|
111
|
-
async def test_specific_search_insights(mock_send_telemetry, mock_insights_retriever, mock_user):
|
|
112
|
-
# Setup
|
|
113
|
-
query = "test query"
|
|
114
|
-
query_type = SearchType.INSIGHTS
|
|
115
|
-
|
|
116
|
-
# Mock the retriever
|
|
117
|
-
mock_retriever = MagicMock()
|
|
118
|
-
mock_retriever.get_completion = AsyncMock()
|
|
119
|
-
mock_retriever.get_completion.return_value = [{"content": "Insight result"}]
|
|
120
|
-
mock_insights_retriever.return_value = mock_retriever
|
|
121
|
-
|
|
122
|
-
# Execute
|
|
123
|
-
results = await specific_search(query_type, query, mock_user)
|
|
124
|
-
|
|
125
|
-
# Verify
|
|
126
|
-
mock_insights_retriever.assert_called_once()
|
|
127
|
-
mock_retriever.get_completion.assert_called_once_with(query)
|
|
128
|
-
mock_send_telemetry.assert_called()
|
|
129
|
-
assert len(results) == 1
|
|
130
|
-
assert results[0]["content"] == "Insight result"
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
@pytest.mark.asyncio
|
|
134
|
-
@patch.object(search_module, "ChunksRetriever")
|
|
135
|
-
@patch.object(search_module, "send_telemetry")
|
|
136
|
-
async def test_specific_search_chunks(mock_send_telemetry, mock_chunks_retriever, mock_user):
|
|
137
|
-
# Setup
|
|
138
|
-
query = "test query"
|
|
139
|
-
query_type = SearchType.CHUNKS
|
|
140
|
-
|
|
141
|
-
# Mock the retriever
|
|
142
|
-
mock_retriever = MagicMock()
|
|
143
|
-
mock_retriever.get_completion = AsyncMock()
|
|
144
|
-
mock_retriever.get_completion.return_value = [{"content": "Chunk result"}]
|
|
145
|
-
mock_chunks_retriever.return_value = mock_retriever
|
|
146
|
-
|
|
147
|
-
# Execute
|
|
148
|
-
results = await specific_search(query_type, query, mock_user)
|
|
149
|
-
|
|
150
|
-
# Verify
|
|
151
|
-
mock_chunks_retriever.assert_called_once()
|
|
152
|
-
mock_retriever.get_completion.assert_called_once_with(query)
|
|
153
|
-
mock_send_telemetry.assert_called()
|
|
154
|
-
assert len(results) == 1
|
|
155
|
-
assert results[0]["content"] == "Chunk result"
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
@pytest.mark.asyncio
|
|
159
|
-
@pytest.mark.parametrize(
|
|
160
|
-
"selected_type, retriever_name, expected_content, top_k",
|
|
161
|
-
[
|
|
162
|
-
(SearchType.RAG_COMPLETION, "CompletionRetriever", "RAG result from lucky search", 10),
|
|
163
|
-
(SearchType.CHUNKS, "ChunksRetriever", "Chunk result from lucky search", 5),
|
|
164
|
-
(SearchType.SUMMARIES, "SummariesRetriever", "Summary from lucky search", 15),
|
|
165
|
-
(SearchType.INSIGHTS, "InsightsRetriever", "Insight result from lucky search", 20),
|
|
166
|
-
],
|
|
167
|
-
)
|
|
168
|
-
@patch.object(search_module, "select_search_type")
|
|
169
|
-
@patch.object(search_module, "send_telemetry")
|
|
170
|
-
async def test_specific_search_feeling_lucky(
|
|
171
|
-
mock_send_telemetry,
|
|
172
|
-
mock_select_search_type,
|
|
173
|
-
selected_type,
|
|
174
|
-
retriever_name,
|
|
175
|
-
expected_content,
|
|
176
|
-
top_k,
|
|
177
|
-
mock_user,
|
|
178
|
-
):
|
|
179
|
-
with patch.object(search_module, retriever_name) as mock_retriever_class:
|
|
180
|
-
# Setup
|
|
181
|
-
query = f"test query for {retriever_name}"
|
|
182
|
-
query_type = SearchType.FEELING_LUCKY
|
|
183
|
-
|
|
184
|
-
# Mock the intelligent search type selection
|
|
185
|
-
mock_select_search_type.return_value = selected_type
|
|
186
|
-
|
|
187
|
-
# Mock the retriever
|
|
188
|
-
mock_retriever_instance = MagicMock()
|
|
189
|
-
mock_retriever_instance.get_completion = AsyncMock(
|
|
190
|
-
return_value=[{"content": expected_content}]
|
|
191
|
-
)
|
|
192
|
-
mock_retriever_class.return_value = mock_retriever_instance
|
|
193
|
-
|
|
194
|
-
# Execute
|
|
195
|
-
results = await specific_search(query_type, query, mock_user, top_k=top_k)
|
|
196
|
-
|
|
197
|
-
# Verify
|
|
198
|
-
mock_select_search_type.assert_called_once_with(query)
|
|
199
|
-
|
|
200
|
-
if retriever_name == "CompletionRetriever":
|
|
201
|
-
mock_retriever_class.assert_called_once_with(
|
|
202
|
-
system_prompt_path="answer_simple_question.txt", top_k=top_k
|
|
203
|
-
)
|
|
204
|
-
else:
|
|
205
|
-
mock_retriever_class.assert_called_once_with(top_k=top_k)
|
|
206
|
-
|
|
207
|
-
mock_retriever_instance.get_completion.assert_called_once_with(query)
|
|
208
|
-
mock_send_telemetry.assert_called()
|
|
209
|
-
assert len(results) == 1
|
|
210
|
-
assert results[0]["content"] == expected_content
|
|
211
|
-
|
|
212
|
-
|
|
213
|
-
@pytest.mark.asyncio
|
|
214
|
-
async def test_specific_search_invalid_type(mock_user):
|
|
215
|
-
# Setup
|
|
216
|
-
query = "test query"
|
|
217
|
-
query_type = "INVALID_TYPE" # Not a valid SearchType
|
|
218
|
-
|
|
219
|
-
# Execute and verify
|
|
220
|
-
with pytest.raises(UnsupportedSearchTypeError) as excinfo:
|
|
221
|
-
await specific_search(query_type, query, mock_user)
|
|
222
|
-
|
|
223
|
-
assert "Unsupported search type" in str(excinfo.value)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|