time_weighted_retriever.py — langchain Source File

Source Code

import datetime
from copy import deepcopy
from typing import Any

from langchain_core.callbacks import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from langchain_core.vectorstores import VectorStore
from pydantic import ConfigDict, Field
from typing_extensions import override


def _get_hours_passed(time: datetime.datetime, ref_time: datetime.datetime) -> float:
    """Get the hours passed between two datetimes."""
    return (time - ref_time).total_seconds() / 3600


class TimeWeightedVectorStoreRetriever(BaseRetriever):
    """Time Weighted Vector Store Retriever.

    Retriever that combines embedding similarity with recency in retrieving values.
    """

    vectorstore: VectorStore
    """The `VectorStore` to store documents and determine salience."""

    search_kwargs: dict = Field(default_factory=lambda: {"k": 100})
    """Keyword arguments to pass to the `VectorStore` similarity search."""

    # TODO: abstract as a queue
    memory_stream: list[Document] = Field(default_factory=list)
    """The memory_stream of documents to search through."""

    decay_rate: float = Field(default=0.01)
    """The exponential decay factor used as `(1.0-decay_rate)**(hrs_passed)`."""

    k: int = 4
    """The maximum number of documents to retrieve in a given call."""

    other_score_keys: list[str] = []
    """Other keys in the metadata to factor into the score, e.g. 'importance'."""

    default_salience: float | None = None
    """The salience to assign memories not retrieved from the vector store.

    None assigns no salience to documents not fetched from the vector store.
    """

    model_config = ConfigDict(
        arbitrary_types_allowed=True,
    )

    def _document_get_date(self, field: str, document: Document) -> datetime.datetime:
        """Return the value of the date field of a document."""
        if field in document.metadata:
            if isinstance(document.metadata[field], float):
                return datetime.datetime.fromtimestamp(document.metadata[field])
// ... (139 more lines)