Koda retriever

KodaRetriever

Bases: BaseRetriever

Custom Hybrid Retriever that dynamically determines the optimal alpha for a given query. An LLM is used to categorize the query and therefore determine the optimal alpha value, as each category has a preset/provided alpha value. It is recommended that you run tests on your corpus of data and queries to determine categories and corresponding alpha values for your use case.

KodaRetriever is built from BaseRetriever, and therefore is a llama-index compatible drop-in replacement for any hybrid retriever.

Auto-routing is NOT enabled without providing an LLM. If no LLM is provided, the default alpha value will be used for all queries and no alpha tuning will be done. Reranking will be done automatically if a reranker is provided. If no matrix is provided, a default matrix is leveraged. (Not recommended for production use)

Parameters:

Name	Type	Description	Default
index	VectorStoreIndex	The index to be used for retrieval	required
llm	LLM	The LLM to be used for auto-routing. Defaults to None.	None
reranker	BaseNodePostprocessor	The reranker to be used for postprocessing. Defaults to None.	None
default_alpha	float	The default alpha value to be used if no LLM is provided. Defaults to .5.	0.5
matrix	dict or AlphaMatrix	The matrix to be used for auto-routing. Defaults to AlphaMatrix(data=DEFAULT_CATEGORIES).	DEFAULT_CATEGORIES
verbose	bool	Whether to log verbose output. Defaults to False.	False
**kwargs		Additional arguments for VectorIndexRetriever	{}

Returns:

Type	Description
	KodaRetriever

Examples:

>>> # Example 1 - provide your own LLM
>>> retriever = KodaRetriever( # woof woof
                    index=vector_index
                    , llm=Settings.llm
                    , verbose=True
                )
>>> results = retriever.retrieve("What is the capital of France?")

>>> # Example 2 - set custom alpha values
>>> matrix_data = { # this is just dummy data, alpha values were randomly chosen
        "positive sentiment": {
            "alpha": .2
            , "description": "Queries expecting a positive answer"
            , "examples": [
                "I love this product"
                , "This is the best product ever"
            ]
        }
        , "negative sentiment": {
            "alpha": .7
            , "description": "Queries expecting a negative answer"
            , "examples": [
                "I hate this product"
                , "This is the worst product ever"
            ]
        }
    }

>>> retriever = KodaRetriever( # woof woof
                    index=vector_index
                    , llm=Settings.llm
                    , matrix=matrix_data
                    , verbose=True
                )
>>> results = retriever.retrieve("What happened on Y2K?")

Source code in llama_index/packs/koda_retriever/base.py

class KodaRetriever(BaseRetriever):
    """
    Custom Hybrid Retriever that dynamically determines the optimal alpha for a given query.
    An LLM is used to categorize the query and therefore determine the optimal alpha value, as each category has a preset/provided alpha value.
    It is recommended that you run tests on your corpus of data and queries to determine categories and corresponding alpha values for your use case.

    KodaRetriever is built from BaseRetriever, and therefore is a llama-index compatible drop-in replacement for any hybrid retriever.

    Auto-routing is NOT enabled without providing an LLM.
    If no LLM is provided, the default alpha value will be used for all queries and no alpha tuning will be done.
    Reranking will be done automatically if a reranker is provided.
    If no matrix is provided, a default matrix is leveraged. (Not recommended for production use)

    Args:
        index (VectorStoreIndex): The index to be used for retrieval
        llm (LLM, optional): The LLM to be used for auto-routing. Defaults to None.
        reranker (BaseNodePostprocessor, optional): The reranker to be used for postprocessing. Defaults to None.
        default_alpha (float, optional): The default alpha value to be used if no LLM is provided. Defaults to .5.
        matrix (dict or AlphaMatrix, optional): The matrix to be used for auto-routing. Defaults to AlphaMatrix(data=DEFAULT_CATEGORIES).
        verbose (bool, optional): Whether to log verbose output. Defaults to False.
        **kwargs: Additional arguments for VectorIndexRetriever

    Returns:
        KodaRetriever

    Examples:
        >>> # Example 1 - provide your own LLM
        >>> retriever = KodaRetriever( # woof woof
                            index=vector_index
                            , llm=Settings.llm
                            , verbose=True
                        )
        >>> results = retriever.retrieve("What is the capital of France?")

        >>> # Example 2 - set custom alpha values
        >>> matrix_data = { # this is just dummy data, alpha values were randomly chosen
                "positive sentiment": {
                    "alpha": .2
                    , "description": "Queries expecting a positive answer"
                    , "examples": [
                        "I love this product"
                        , "This is the best product ever"
                    ]
                }
                , "negative sentiment": {
                    "alpha": .7
                    , "description": "Queries expecting a negative answer"
                    , "examples": [
                        "I hate this product"
                        , "This is the worst product ever"
                    ]
                }
            }

        >>> retriever = KodaRetriever( # woof woof
                            index=vector_index
                            , llm=Settings.llm
                            , matrix=matrix_data
                            , verbose=True
                        )
        >>> results = retriever.retrieve("What happened on Y2K?")

    """

    def __init__(
        self,
        index: VectorStoreIndex,
        llm: Optional[LLM] = None,  # if I could, I'd default to
        reranker: Optional[BaseNodePostprocessor] = None,
        default_alpha: float = 0.5,
        matrix: dict or AlphaMatrix = DEFAULT_CATEGORIES,  # type: ignore
        verbose: bool = False,
        **kwargs,  # kwargs for VectorIndexRetriever
    ):
        super().__init__()

        self.index = index
        self.retriever = VectorIndexRetriever(
            index=index,
            vector_store_query_mode=VectorStoreQueryMode.HYBRID,
            alpha=default_alpha,
            **kwargs,  # filters, etc, added here
        )
        self.default_alpha = default_alpha
        self.reranker = reranker
        self.llm = llm
        self.verbose = verbose

        if isinstance(matrix, dict):
            matrix = AlphaMatrix(data=matrix)

        self.matrix = matrix

    def categorize(self, query: str) -> str:
        """Categorizes a query using the provided LLM and matrix. If no LLM is provided, the default alpha value will be used."""
        if not self.llm:
            err = "LLM is required for auto-routing. During instantiation, please provide an LLM or use direct routing."
            raise TypeError(err)

        prompt = CATEGORIZER_PROMPT.format(
            question=query, category_info=self.matrix.get_all_category_info()
        )

        response = str(self.llm.complete(prompt))  # type: ignore

        if response not in self.matrix.get_categories():
            raise ValueError(
                f"LLM classified question in a category that is not registered. {response} not in {self.matrix.get_categories()}"
            )

        return response

    async def a_categorize(self, query: str) -> str:
        """(async) Categorizes a query using the provided LLM and matrix. If no LLM is provided, the default alpha value will be used."""
        if not self.llm:
            err = "LLM is required for auto-routing. During instantiation, please provide an LLM or use direct routing."
            raise TypeError(err)

        prompt = CATEGORIZER_PROMPT.format(
            question=query, category_info=self.matrix.get_all_category_info()
        )

        response = await self.llm.acomplete(prompt)  # type: ignore
        response = str(response)

        if response not in self.matrix.get_categories():
            raise ValueError(
                f"LLM classified question in a category that is not registered. {response!s} not in {self.matrix.get_categories()}"
            )

        return response

    def category_retrieve(self, category: str, query: QueryType) -> List[NodeWithScore]:
        """Updates the alpha and retrieves results for a query using the provided category and query. If no LLM is provided, the default alpha value will be used."""
        alpha = self.matrix.get_alpha(category)
        self.retriever._alpha = (
            alpha  # updates alpha according to classification of query
        )

        return self.retriever.retrieve(str_or_query_bundle=query)

    async def a_category_retrieve(
        self, category: str, query: QueryType
    ) -> List[NodeWithScore]:
        """(async) Updates the alpha and retrieves results for a query using the provided category and query. If no LLM is provided, the default alpha value will be used."""
        alpha = self.matrix.get_alpha(category)
        self.retriever._alpha = (
            alpha  # updates alpha according to classification of query
        )

        return await self.retriever.aretrieve(str_or_query_bundle=query)

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """llama-index compatible retrieve method that auto-determines the optimal alpha for a query and then retrieves results for a query."""
        if not self.llm:
            warning = f"LLM is not provided, skipping route categorization. Default alpha of {self.default_alpha} will be used."
            logging.warning(warning)

            results = self.retriever.retrieve(query_bundle)
        else:
            category = self.categorize(query=query_bundle.query_str)  # type: ignore
            results = self.category_retrieve(category, query_bundle)
            if self.verbose:
                logging.info(
                    f"Query categorized as {category} with alpha of {self.matrix.get_alpha(category)}"
                )

        if self.reranker:
            if self.verbose:
                logging.info("Reranking results")
            results = self.reranker.postprocess_nodes(
                nodes=results, query_bundle=query_bundle
            )

        return results

    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """(async) llama-index compatible retrieve method that auto-determines the optimal alpha for a query and then retrieves results for a query."""
        if not self.llm:
            warning = f"LLM is not provided, skipping route categorization. Default alpha of {self.default_alpha} will be used."
            logging.warning(warning)

            results = await self.retriever.aretrieve(query_bundle)
        else:
            category = await self.a_categorize(query_bundle.query_str)  # type: ignore
            results = await self.a_category_retrieve(category, query_bundle)
            if self.verbose:
                logging.info(
                    f"Query categorized as {category} with alpha of {self.matrix.get_alpha(category)}"
                )

        if self.reranker:
            if self.verbose:
                logging.info("Reranking results")
            results = self.reranker.postprocess_nodes(
                nodes=results, query_bundle=query_bundle
            )

        return results

categorize

categorize(query: str) -> str

Categorizes a query using the provided LLM and matrix. If no LLM is provided, the default alpha value will be used.

Source code in llama_index/packs/koda_retriever/base.py

def categorize(self, query: str) -> str:
    """Categorizes a query using the provided LLM and matrix. If no LLM is provided, the default alpha value will be used."""
    if not self.llm:
        err = "LLM is required for auto-routing. During instantiation, please provide an LLM or use direct routing."
        raise TypeError(err)

    prompt = CATEGORIZER_PROMPT.format(
        question=query, category_info=self.matrix.get_all_category_info()
    )

    response = str(self.llm.complete(prompt))  # type: ignore

    if response not in self.matrix.get_categories():
        raise ValueError(
            f"LLM classified question in a category that is not registered. {response} not in {self.matrix.get_categories()}"
        )

    return response

a_categorize async

a_categorize(query: str) -> str

(async) Categorizes a query using the provided LLM and matrix. If no LLM is provided, the default alpha value will be used.

Source code in llama_index/packs/koda_retriever/base.py

async def a_categorize(self, query: str) -> str:
    """(async) Categorizes a query using the provided LLM and matrix. If no LLM is provided, the default alpha value will be used."""
    if not self.llm:
        err = "LLM is required for auto-routing. During instantiation, please provide an LLM or use direct routing."
        raise TypeError(err)

    prompt = CATEGORIZER_PROMPT.format(
        question=query, category_info=self.matrix.get_all_category_info()
    )

    response = await self.llm.acomplete(prompt)  # type: ignore
    response = str(response)

    if response not in self.matrix.get_categories():
        raise ValueError(
            f"LLM classified question in a category that is not registered. {response!s} not in {self.matrix.get_categories()}"
        )

    return response

category_retrieve

category_retrieve(category: str, query: QueryType) -> List[NodeWithScore]

Updates the alpha and retrieves results for a query using the provided category and query. If no LLM is provided, the default alpha value will be used.

Source code in llama_index/packs/koda_retriever/base.py

def category_retrieve(self, category: str, query: QueryType) -> List[NodeWithScore]:
    """Updates the alpha and retrieves results for a query using the provided category and query. If no LLM is provided, the default alpha value will be used."""
    alpha = self.matrix.get_alpha(category)
    self.retriever._alpha = (
        alpha  # updates alpha according to classification of query
    )

    return self.retriever.retrieve(str_or_query_bundle=query)

a_category_retrieve async

a_category_retrieve(category: str, query: QueryType) -> List[NodeWithScore]

(async) Updates the alpha and retrieves results for a query using the provided category and query. If no LLM is provided, the default alpha value will be used.

Source code in llama_index/packs/koda_retriever/base.py

async def a_category_retrieve(
    self, category: str, query: QueryType
) -> List[NodeWithScore]:
    """(async) Updates the alpha and retrieves results for a query using the provided category and query. If no LLM is provided, the default alpha value will be used."""
    alpha = self.matrix.get_alpha(category)
    self.retriever._alpha = (
        alpha  # updates alpha according to classification of query
    )

    return await self.retriever.aretrieve(str_or_query_bundle=query)

AlphaMatrix

Bases: BaseModel

This class is not necessary to understand to use a KodaRetriever - as it will be automatically instantiated if a dictionary is provided.

Pydantic class to enforce the required fields for a KodaRetriever Its best to just instantiate this using a dictionary, don't both trying to instantiate by declaring any AlphaCategory objects.

Example

data = { "normal query": { # examples is not required if you aren't using few-shot auto-routing "alpha": .5 , "description": "This is a normal query" # desc is not required if you aren't using few-shot auto-routing , "examples": ["This is a normal query", "Another normal query"] } } matrix = AlphaMatrix(data=data) # arg must be named matrix for the retriever to use it

Source code in llama_index/packs/koda_retriever/matrix.py

class AlphaMatrix(BaseModel):
    """
    This class is not necessary to understand to use a KodaRetriever - as it will be automatically instantiated if a dictionary is provided.

    Pydantic class to enforce the required fields for a KodaRetriever
    Its best to just instantiate this using a dictionary, don't both trying to instantiate by declaring any AlphaCategory objects.

    Example:
        >>> data = {
                "normal query": { # examples is not required if you aren't using few-shot auto-routing
                    "alpha": .5
                    , "description": "This is a normal query" # desc is not required if you aren't using few-shot auto-routing
                    , "examples": ["This is a normal query", "Another normal query"]
                }
            }
        >>> matrix = AlphaMatrix(data=data) # arg must be named matrix for the retriever to use it

    """

    class AlphaCategory(BaseModel):
        """
        Subclass to enforce the required fields for a category in the AlphaMatrix - necessary for nesting in the AlphaMatrix class
        You should not have to really touch this, as it is only used for type checking and validation.
        """

        alpha: float
        description: Optional[str] = (
            None  # optional if providing a custom LLM, its presumed this was part of your training data for the custom model
        )
        examples: Optional[List[str]] = (
            None  # if not providing a custom model, this is required
        )

    data: Dict[str, AlphaCategory]

    def get_alpha(self, category: str) -> float:
        """Simple helper function to get the alpha value for a given category."""
        if category not in self.data:
            err = f"Provided category '{category}' cannot be found"
            raise ValueError(err)

        return self.data.get(category).alpha  # type: ignore

    def get_examples(self, category: str) -> List[str]:
        """Simple helper function to get the examples for a given category."""
        if category not in self.data:
            err = f"Provided category '{category}' cannot be found"
            raise ValueError(err)

        return self.data.get(category).examples  # type: ignore

    def get_description(self, category: str) -> str:
        """Simple helper function to get the description for a given category."""
        if category not in self.data:
            err = f"Provided category '{category}' cannot be found"
            raise ValueError(err)

        return self.data.get(category).description  # type: ignore

    def get_categories(self) -> list:
        """Simple helper function to get the categories for a given category."""
        return list(self.data.keys())

    def format_category(self, category: str) -> str:
        """Simple helper function to format the category information for a given category."""
        if category not in self.data:
            err = f"Provided category '{category}' cannot be found"
            raise ValueError(err)

        description = self.get_description(category)
        examples = self.get_examples(category)

        category_info = f"""
        - {category}:
            description: {description}
        """.strip()

        if examples:
            examples = "; ".join(examples)
            example_info = f"""
            examples:
                {examples}
            """
            category_info = f"{category_info}\n{example_info}"

        return category_info

    def get_all_category_info(self) -> str:
        """Simple helper function to get the category information for all categories."""
        categories = []
        for category in self.get_categories():
            category_info = self.format_category(category)
            categories.append(category_info)
        return "\n".join(categories)

AlphaCategory

Bases: BaseModel

Subclass to enforce the required fields for a category in the AlphaMatrix - necessary for nesting in the AlphaMatrix class You should not have to really touch this, as it is only used for type checking and validation.

Source code in llama_index/packs/koda_retriever/matrix.py

class AlphaCategory(BaseModel):
    """
    Subclass to enforce the required fields for a category in the AlphaMatrix - necessary for nesting in the AlphaMatrix class
    You should not have to really touch this, as it is only used for type checking and validation.
    """

    alpha: float
    description: Optional[str] = (
        None  # optional if providing a custom LLM, its presumed this was part of your training data for the custom model
    )
    examples: Optional[List[str]] = (
        None  # if not providing a custom model, this is required
    )

get_alpha

get_alpha(category: str) -> float

Simple helper function to get the alpha value for a given category.

Source code in llama_index/packs/koda_retriever/matrix.py

def get_alpha(self, category: str) -> float:
    """Simple helper function to get the alpha value for a given category."""
    if category not in self.data:
        err = f"Provided category '{category}' cannot be found"
        raise ValueError(err)

    return self.data.get(category).alpha  # type: ignore

get_examples

get_examples(category: str) -> List[str]

Simple helper function to get the examples for a given category.

Source code in llama_index/packs/koda_retriever/matrix.py

def get_examples(self, category: str) -> List[str]:
    """Simple helper function to get the examples for a given category."""
    if category not in self.data:
        err = f"Provided category '{category}' cannot be found"
        raise ValueError(err)

    return self.data.get(category).examples  # type: ignore

get_description

get_description(category: str) -> str

Simple helper function to get the description for a given category.

Source code in llama_index/packs/koda_retriever/matrix.py

def get_description(self, category: str) -> str:
    """Simple helper function to get the description for a given category."""
    if category not in self.data:
        err = f"Provided category '{category}' cannot be found"
        raise ValueError(err)

    return self.data.get(category).description  # type: ignore

get_categories

get_categories() -> list

Simple helper function to get the categories for a given category.

Source code in llama_index/packs/koda_retriever/matrix.py

def get_categories(self) -> list:
    """Simple helper function to get the categories for a given category."""
    return list(self.data.keys())

format_category

format_category(category: str) -> str

Simple helper function to format the category information for a given category.

Source code in llama_index/packs/koda_retriever/matrix.py

def format_category(self, category: str) -> str:
    """Simple helper function to format the category information for a given category."""
    if category not in self.data:
        err = f"Provided category '{category}' cannot be found"
        raise ValueError(err)

    description = self.get_description(category)
    examples = self.get_examples(category)

    category_info = f"""
    - {category}:
        description: {description}
    """.strip()

    if examples:
        examples = "; ".join(examples)
        example_info = f"""
        examples:
            {examples}
        """
        category_info = f"{category_info}\n{example_info}"

    return category_info

get_all_category_info

get_all_category_info() -> str

Simple helper function to get the category information for all categories.

Source code in llama_index/packs/koda_retriever/matrix.py

def get_all_category_info(self) -> str:
    """Simple helper function to get the category information for all categories."""
    categories = []
    for category in self.get_categories():
        category_info = self.format_category(category)
        categories.append(category_info)
    return "\n".join(categories)

options: members: - KodaRetrieverPack