chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,348 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
import logging
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Callable, Sequence
|
||||
from copy import deepcopy
|
||||
from typing import Any, Final, Literal, TypeVar, overload
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from semantic_kernel.data._shared import (
|
||||
DEFAULT_FUNCTION_NAME,
|
||||
DEFAULT_PARAMETER_METADATA,
|
||||
DEFAULT_RETURN_PARAMETER_METADATA,
|
||||
DynamicFilterFunction,
|
||||
KernelSearchResults,
|
||||
SearchOptions,
|
||||
create_options,
|
||||
default_dynamic_filter_function,
|
||||
)
|
||||
from semantic_kernel.exceptions import TextSearchException
|
||||
from semantic_kernel.functions.kernel_function import KernelFunction
|
||||
from semantic_kernel.functions.kernel_function_decorator import kernel_function
|
||||
from semantic_kernel.functions.kernel_function_from_method import KernelFunctionFromMethod
|
||||
from semantic_kernel.functions.kernel_parameter_metadata import KernelParameterMetadata
|
||||
from semantic_kernel.kernel_pydantic import KernelBaseModel
|
||||
from semantic_kernel.kernel_types import OptionalOneOrList
|
||||
from semantic_kernel.utils.feature_stage_decorator import release_candidate
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TSearchOptions = TypeVar("TSearchOptions", bound="SearchOptions")
|
||||
|
||||
DEFAULT_DESCRIPTION: Final[str] = (
|
||||
"Perform a search for content related to the specified query and return string results"
|
||||
)
|
||||
|
||||
# region: Results
|
||||
|
||||
|
||||
@release_candidate
|
||||
class TextSearchResult(KernelBaseModel):
|
||||
"""The result of a text search."""
|
||||
|
||||
name: str | None = None
|
||||
value: str | None = None
|
||||
link: str | None = None
|
||||
|
||||
|
||||
TSearchResult = TypeVar("TSearchResult")
|
||||
|
||||
|
||||
@release_candidate
|
||||
class TextSearch:
|
||||
"""The base class for all text searchers."""
|
||||
|
||||
@property
|
||||
def options_class(self) -> type["SearchOptions"]:
|
||||
"""The options class for the search."""
|
||||
return SearchOptions
|
||||
|
||||
# region: Public methods
|
||||
|
||||
@overload
|
||||
def create_search_function(
|
||||
self,
|
||||
function_name: str = DEFAULT_FUNCTION_NAME,
|
||||
description: str = DEFAULT_DESCRIPTION,
|
||||
*,
|
||||
output_type: Literal["str"] = "str",
|
||||
parameters: list[KernelParameterMetadata] | None = None,
|
||||
return_parameter: KernelParameterMetadata | None = None,
|
||||
filter: OptionalOneOrList[Callable | str] = None,
|
||||
top: int = 5,
|
||||
skip: int = 0,
|
||||
include_total_count: bool = False,
|
||||
filter_update_function: DynamicFilterFunction | None = None,
|
||||
string_mapper: Callable[[TSearchResult], str] | None = None,
|
||||
) -> KernelFunction:
|
||||
"""Create a kernel function from a search function.
|
||||
|
||||
Args:
|
||||
output_type: The type of the output, default is "str".
|
||||
function_name: The name of the function, to be used in the kernel, default is "search".
|
||||
description: The description of the function, a default is provided.
|
||||
parameters: The parameters for the function, a list of KernelParameterMetadata.
|
||||
return_parameter: The return parameter for the function.
|
||||
filter: The filter to use for the search.
|
||||
top: The number of results to return.
|
||||
skip: The number of results to skip.
|
||||
include_total_count: Whether to include the total count of results.
|
||||
filter_update_function: A function to update the search filters.
|
||||
The function should return the updated filter.
|
||||
The default function uses the parameters and the kwargs to update the options.
|
||||
Adding equal to filters to the options for all parameters that are not "query".
|
||||
As well as adding equal to filters for parameters that have a default value.
|
||||
string_mapper: The function to map the search results. (the inner part of the KernelSearchResults type,
|
||||
related to which search type you are using) to strings.
|
||||
|
||||
Returns:
|
||||
KernelFunction: The kernel function.
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
@overload
|
||||
def create_search_function(
|
||||
self,
|
||||
function_name: str = DEFAULT_FUNCTION_NAME,
|
||||
description: str = DEFAULT_DESCRIPTION,
|
||||
*,
|
||||
output_type: Literal["TextSearchResult"],
|
||||
parameters: list[KernelParameterMetadata] | None = None,
|
||||
return_parameter: KernelParameterMetadata | None = None,
|
||||
filter: OptionalOneOrList[Callable | str] = None,
|
||||
top: int = 5,
|
||||
skip: int = 0,
|
||||
include_total_count: bool = False,
|
||||
filter_update_function: DynamicFilterFunction | None = None,
|
||||
) -> KernelFunction:
|
||||
"""Create a kernel function from a search function.
|
||||
|
||||
Args:
|
||||
output_type: The type of the output, in this case TextSearchResult.
|
||||
function_name: The name of the function, to be used in the kernel, default is "search".
|
||||
description: The description of the function, a default is provided.
|
||||
parameters: The parameters for the function, a list of KernelParameterMetadata.
|
||||
return_parameter: The return parameter for the function.
|
||||
filter: The filter to use for the search.
|
||||
top: The number of results to return.
|
||||
skip: The number of results to skip.
|
||||
include_total_count: Whether to include the total count of results.
|
||||
filter_update_function: A function to update the search filters.
|
||||
The function should return the updated filter.
|
||||
The default function uses the parameters and the kwargs to update the options.
|
||||
Adding equal to filters to the options for all parameters that are not "query".
|
||||
As well as adding equal to filters for parameters that have a default value.
|
||||
string_mapper: The function to map the TextSearchResult to strings.
|
||||
for instance taking the value out of the results and just returning that,
|
||||
otherwise a json-like string is returned.
|
||||
|
||||
Returns:
|
||||
KernelFunction: The kernel function.
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
@overload
|
||||
def create_search_function(
|
||||
self,
|
||||
function_name: str = DEFAULT_FUNCTION_NAME,
|
||||
description: str = DEFAULT_DESCRIPTION,
|
||||
*,
|
||||
output_type: Literal["Any"],
|
||||
parameters: list[KernelParameterMetadata] | None = None,
|
||||
return_parameter: KernelParameterMetadata | None = None,
|
||||
filter: OptionalOneOrList[Callable | str] = None,
|
||||
top: int = 5,
|
||||
skip: int = 0,
|
||||
include_total_count: bool = False,
|
||||
filter_update_function: DynamicFilterFunction | None = None,
|
||||
) -> KernelFunction:
|
||||
"""Create a kernel function from a search function.
|
||||
|
||||
Args:
|
||||
function_name: The name of the function, to be used in the kernel, default is "search".
|
||||
description: The description of the function, a default is provided.
|
||||
output_type: The type of the output, in this case Any.
|
||||
Any means that the results from the store are used directly.
|
||||
The string_mapper can then be used to extract certain fields.
|
||||
parameters: The parameters for the function, a list of KernelParameterMetadata.
|
||||
return_parameter: The return parameter for the function.
|
||||
filter: The filter to use for the search.
|
||||
top: The number of results to return.
|
||||
skip: The number of results to skip.
|
||||
include_total_count: Whether to include the total count of results.
|
||||
filter_update_function: A function to update the search filters.
|
||||
The function should return the updated filter.
|
||||
The default function uses the parameters and the kwargs to update the options.
|
||||
Adding equal to filters to the options for all parameters that are not "query".
|
||||
As well as adding equal to filters for parameters that have a default value.
|
||||
string_mapper: The function to map the raw search results to strings.
|
||||
When using this from a vector store, your results are of type
|
||||
VectorSearchResult[TModel],
|
||||
so the string_mapper can be used to extract the fields you want from the result.
|
||||
The default is to use the model_dump_json method of the result, which will return a json-like string.
|
||||
|
||||
Returns:
|
||||
KernelFunction: The kernel function.
|
||||
"""
|
||||
...
|
||||
|
||||
def create_search_function(
|
||||
self,
|
||||
function_name=DEFAULT_FUNCTION_NAME,
|
||||
description=DEFAULT_DESCRIPTION,
|
||||
*,
|
||||
output_type="str",
|
||||
parameters=None,
|
||||
return_parameter=None,
|
||||
filter=None,
|
||||
top=5,
|
||||
skip=0,
|
||||
include_total_count=False,
|
||||
filter_update_function=None,
|
||||
string_mapper=None,
|
||||
) -> KernelFunction:
|
||||
"""Create a kernel function from a search function."""
|
||||
options = SearchOptions(
|
||||
filter=filter,
|
||||
skip=skip,
|
||||
top=top,
|
||||
include_total_count=include_total_count,
|
||||
)
|
||||
match output_type:
|
||||
case "str":
|
||||
return self._create_kernel_function(
|
||||
output_type=str,
|
||||
options=options,
|
||||
parameters=parameters,
|
||||
filter_update_function=filter_update_function,
|
||||
return_parameter=return_parameter,
|
||||
function_name=function_name,
|
||||
description=description,
|
||||
string_mapper=string_mapper,
|
||||
)
|
||||
case "TextSearchResult":
|
||||
return self._create_kernel_function(
|
||||
output_type=TextSearchResult,
|
||||
options=options,
|
||||
parameters=parameters,
|
||||
filter_update_function=filter_update_function,
|
||||
return_parameter=return_parameter,
|
||||
function_name=function_name,
|
||||
description=description,
|
||||
string_mapper=string_mapper,
|
||||
)
|
||||
case "Any":
|
||||
return self._create_kernel_function(
|
||||
output_type="Any",
|
||||
options=options,
|
||||
parameters=parameters,
|
||||
filter_update_function=filter_update_function,
|
||||
return_parameter=return_parameter,
|
||||
function_name=function_name,
|
||||
description=description,
|
||||
string_mapper=string_mapper,
|
||||
)
|
||||
case _:
|
||||
raise TextSearchException(
|
||||
f"Unknown output type: {output_type}. Must be 'str', 'TextSearchResult', or 'Any'."
|
||||
)
|
||||
|
||||
# endregion
|
||||
# region: Private methods
|
||||
|
||||
def _create_kernel_function(
|
||||
self,
|
||||
output_type: type[str] | type[TSearchResult] | Literal["Any"] = str,
|
||||
options: SearchOptions | None = None,
|
||||
parameters: list[KernelParameterMetadata] | None = None,
|
||||
filter_update_function: DynamicFilterFunction | None = None,
|
||||
return_parameter: KernelParameterMetadata | None = None,
|
||||
function_name: str = DEFAULT_FUNCTION_NAME,
|
||||
description: str = DEFAULT_DESCRIPTION,
|
||||
string_mapper: Callable[[TSearchResult], str] | None = None,
|
||||
) -> KernelFunction:
|
||||
"""Create a kernel function from a search function."""
|
||||
update_func = filter_update_function or default_dynamic_filter_function
|
||||
|
||||
@kernel_function(name=function_name, description=description)
|
||||
async def search_wrapper(**kwargs: Any) -> Sequence[str]:
|
||||
query = kwargs.pop("query", "")
|
||||
try:
|
||||
inner_options = create_options(SearchOptions, deepcopy(options), **kwargs)
|
||||
except ValidationError:
|
||||
# this usually only happens when the kwargs are invalid, so blank options in this case.
|
||||
inner_options = SearchOptions()
|
||||
inner_options.filter = update_func(filter=inner_options.filter, parameters=parameters, **kwargs)
|
||||
try:
|
||||
results = await self.search(
|
||||
query=query,
|
||||
output_type=output_type,
|
||||
**inner_options.model_dump(exclude_none=True, exclude_defaults=True, exclude_unset=True),
|
||||
)
|
||||
except Exception as e:
|
||||
msg = f"Exception in search function: {e}"
|
||||
logger.error(msg)
|
||||
raise TextSearchException(msg) from e
|
||||
return await self._map_results(results, string_mapper)
|
||||
|
||||
return KernelFunctionFromMethod(
|
||||
method=search_wrapper,
|
||||
parameters=DEFAULT_PARAMETER_METADATA if parameters is None else parameters,
|
||||
return_parameter=return_parameter or DEFAULT_RETURN_PARAMETER_METADATA,
|
||||
)
|
||||
|
||||
async def _map_results(
|
||||
self,
|
||||
results: KernelSearchResults[TSearchResult],
|
||||
string_mapper: Callable[[TSearchResult], str] | None = None,
|
||||
) -> list[str]:
|
||||
"""Map search results to strings."""
|
||||
if string_mapper:
|
||||
return [string_mapper(result) async for result in results.results]
|
||||
return [self._default_map_to_string(result) async for result in results.results]
|
||||
|
||||
@staticmethod
|
||||
def _default_map_to_string(result: BaseModel | object) -> str:
|
||||
"""Default mapping function for text search results."""
|
||||
if isinstance(result, BaseModel):
|
||||
return result.model_dump_json()
|
||||
return result if isinstance(result, str) else json.dumps(result)
|
||||
|
||||
# region: Abstract methods
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
query: str,
|
||||
output_type: type[str] | type[TSearchResult] | Literal["Any"] = str,
|
||||
**kwargs: Any,
|
||||
) -> "KernelSearchResults[TSearchResult]":
|
||||
"""Search for text, returning a KernelSearchResult with a list of strings.
|
||||
|
||||
Args:
|
||||
query: The query to search for.
|
||||
output_type: The type of the output, default is str.
|
||||
Can also be TextSearchResult or Any.
|
||||
**kwargs: Additional keyword arguments to pass to the search function.
|
||||
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_DESCRIPTION",
|
||||
"DEFAULT_FUNCTION_NAME",
|
||||
"DEFAULT_PARAMETER_METADATA",
|
||||
"DEFAULT_RETURN_PARAMETER_METADATA",
|
||||
"DynamicFilterFunction",
|
||||
"KernelSearchResults",
|
||||
"TextSearch",
|
||||
"TextSearchResult",
|
||||
"create_options",
|
||||
"default_dynamic_filter_function",
|
||||
]
|
||||
Reference in New Issue
Block a user