d7c20c1f59
* WIP * WIP * WIP * community search * WIP * WIP * integration tested * tests * tests * mypy * mypy * format
84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
"""
|
|
Copyright 2024, Zep Software, Inc.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
from graphiti_core.search.search_config import (
|
|
CommunityReranker,
|
|
CommunitySearchConfig,
|
|
CommunitySearchMethod,
|
|
EdgeReranker,
|
|
EdgeSearchConfig,
|
|
EdgeSearchMethod,
|
|
NodeReranker,
|
|
NodeSearchConfig,
|
|
NodeSearchMethod,
|
|
SearchConfig,
|
|
)
|
|
|
|
# Performs a hybrid search with rrf reranking over edges, nodes, and communities
|
|
COMBINED_HYBRID_SEARCH_RRF = SearchConfig(
|
|
edge_config=EdgeSearchConfig(
|
|
search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
|
|
reranker=EdgeReranker.rrf,
|
|
),
|
|
node_config=NodeSearchConfig(
|
|
search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
|
|
reranker=NodeReranker.rrf,
|
|
),
|
|
community_config=CommunitySearchConfig(
|
|
search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
|
|
reranker=CommunityReranker.rrf,
|
|
),
|
|
)
|
|
|
|
# performs a hybrid search over edges with rrf reranking
|
|
EDGE_HYBRID_SEARCH_RRF = SearchConfig(
|
|
edge_config=EdgeSearchConfig(
|
|
search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
|
|
reranker=EdgeReranker.rrf,
|
|
)
|
|
)
|
|
|
|
# performs a hybrid search over edges with node distance reranking
|
|
EDGE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
|
|
edge_config=EdgeSearchConfig(
|
|
search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
|
|
reranker=EdgeReranker.node_distance,
|
|
)
|
|
)
|
|
|
|
# performs a hybrid search over nodes with rrf reranking
|
|
NODE_HYBRID_SEARCH_RRF = SearchConfig(
|
|
node_config=NodeSearchConfig(
|
|
search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
|
|
reranker=NodeReranker.rrf,
|
|
)
|
|
)
|
|
|
|
# performs a hybrid search over nodes with node distance reranking
|
|
NODE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
|
|
node_config=NodeSearchConfig(
|
|
search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
|
|
reranker=NodeReranker.node_distance,
|
|
)
|
|
)
|
|
|
|
# performs a hybrid search over communities with rrf reranking
|
|
COMMUNITY_HYBRID_SEARCH_RRF = SearchConfig(
|
|
community_config=CommunitySearchConfig(
|
|
search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
|
|
reranker=CommunityReranker.rrf,
|
|
)
|
|
)
|