rebase

Author: liyin2015

adalflow/adalflow/components/retriever/__init__.py

@@ -22,12 +22,18 @@
     OptionalPackages.SQLALCHEMY,
 )
 
+QdrantRetriever = LazyImport(
+    "adalflow.components.retriever.qdrant_retriever.QdrantRetriever",
+    OptionalPackages.QDRANT,
+)
+
 __all__ = [
     "BM25Retriever",
     "LLMRetriever",
     "FAISSRetriever",
     "RerankerRetriever",
     "PostgresRetriever",
+    "QdrantRetriever",
     "split_text_by_word_fn",
     "split_text_by_word_fn_then_lower_tokenized",
 ]

adalflow/adalflow/components/retriever/qdrant_retriever.py

@@ -0,0 +1,159 @@
+"""Leverage a Qdrant collection to retrieve documents."""
+
+from typing import List, Optional, Any
+from qdrant_client import QdrantClient, models
+
+from adalflow.core.retriever import (
+    Retriever,
+)
+from adalflow.core.embedder import Embedder
+
+from adalflow.core.types import (
+    RetrieverOutput,
+    RetrieverStrQueryType,
+    RetrieverStrQueriesType,
+    Document,
+)
+
+
+class QdrantRetriever(Retriever[Any, RetrieverStrQueryType]):
+    __doc__ = r"""Use a Qdrant collection to retrieve documents.
+
+    Args:
+        collection_name (str): the collection name in Qdrant.
+        client (QdrantClient): An instance of qdrant_client.QdrantClient.
+        embedder (Embedder): An instance of Embedder.
+        top_k (Optional[int], optional): top k documents to fetch. Defaults to 10.
+        vector_name (Optional[str], optional): the name of the vector in the collection. Defaults to None.
+        text_key (str, optional): the key in the payload that contains the text. Defaults to "text".
+        metadata_key (str, optional): the key in the payload that contains the metadata. Defaults to "meta_data".
+        filter (Optional[models.Filter], optional): the filter to apply to the query. Defaults to None.
+
+    References:
+    [1] Qdrant: https://qdrant.tech/
+    [2] Documentation: https://qdrant.tech/documentation/
+    """
+
+    def __init__(
+        self,
+        collection_name: str,
+        client: QdrantClient,
+        embedder: Embedder,
+        top_k: Optional[int] = 10,
+        vector_name: Optional[str] = None,
+        text_key: str = "text",
+        metadata_key: str = "meta_data",
+        filter: Optional[models.Filter] = None,
+    ):
+        super().__init__()
+        self._top_k = top_k
+        self._collection_name = collection_name
+        self._client = client
+        self._embedder = embedder
+        self._text_key = text_key
+        self._metadata_key = metadata_key
+        self._filter = filter
+
+        self._vector_name = vector_name or self._get_first_vector_name()
+
+    def reset_index(self):
+        if self._client.collection_exists(self._collection_name):
+            self._client.delete_collection(self._collection_name)
+
+    def call(
+        self,
+        input: RetrieverStrQueriesType,
+        top_k: Optional[int] = None,
+        **kwargs,
+    ) -> List[RetrieverOutput]:
+        top_k = top_k or self._top_k
+        queries: List[str] = input if isinstance(input, list) else [input]
+
+        queries_embeddings = self._embedder(queries)
+
+        query_requests: List[models.QueryRequest] = []
+        for idx, query in enumerate(queries):
+            query_embedding = queries_embeddings.data[idx].embedding
+            query_requests.append(
+                models.QueryRequest(
+                    query=query_embedding,
+                    limit=top_k,
+                    using=self._vector_name,
+                    with_payload=True,
+                    with_vector=True,
+                    filter=self._filter,
+                    **kwargs,
+                )
+            )
+
+        results = self._client.query_batch_points(
+            self._collection_name, requests=query_requests
+        )
+        retrieved_outputs: List[RetrieverOutput] = []
+        for result in results:
+            out = self._points_to_output(
+                result.points,
+                query,
+                self._text_key,
+                self._metadata_key,
+                self._vector_name,
+            )
+            retrieved_outputs.append(out)
+
+        return retrieved_outputs
+
+    def _get_first_vector_name(self) -> Optional[str]:
+        vectors = self._client.get_collection(
+            self._collection_name
+        ).config.params.vectors
+
+        if not isinstance(vectors, dict):
+            # The collection only has the default, unnamed vector
+            return None
+
+        first_vector_name = list(vectors.keys())[0]
+
+        # The collection has multiple vectors. Could also include the falsy unnamed vector - Empty string("")
+        return first_vector_name or None
+
+    @classmethod
+    def _points_to_output(
+        cls,
+        points: List[models.ScoredPoint],
+        query: str,
+        text_key: str,
+        metadata_key: str,
+        vector_name: Optional[str],
+    ) -> RetrieverOutput:
+        doc_indices = [point.id for point in points]
+        doc_scores = [point.score for point in points]
+        documents = [
+            cls._doc_from_point(point, text_key, metadata_key, vector_name)
+            for point in points
+        ]
+        return RetrieverOutput(
+            doc_indices=doc_indices,
+            doc_scores=doc_scores,
+            query=query,
+            documents=documents,
+        )
+
+    @classmethod
+    def _doc_from_point(
+        cls,
+        point: models.ScoredPoint,
+        text_key: str,
+        metadata_key: str,
+        vector_name: Optional[str] = None,
+    ) -> Document:
+        vector = point.vector
+        if isinstance(vector, dict):
+            vector = vector[vector_name]
+
+        payload = point.payload.copy()
+        return Document(
+            id=point.id,
+            text=payload.get(text_key, ""),
+            meta_data=payload.get(metadata_key, {}),
+            vector=vector,
+        )

adalflow/adalflow/core/generator.py

@@ -810,7 +810,12 @@ async def acall(
                 output = GeneratorOutput(raw_response=str(completion), error=str(e))
 
         log.info(f"output: {output}")
-        self._run_callbacks(output, input=api_kwargs)
+        self._run_callbacks(
+            output,
+            input=api_kwargs,
+            prompt_kwargs=prompt_kwargs,
+            model_kwargs=model_kwargs,
+        )
         return output
 
     def __call__(self, *args, **kwargs) -> Union[GeneratorOutputType, Any]:

adalflow/adalflow/utils/lazy_import.py

@@ -51,6 +51,10 @@ class OptionalPackages(Enum):
         "datasets",
         "Please install datasets with: pip install datasets",
     )
+    QDRANT = (
+        "qdrant-client",
+        "Please install qdrant-client with: pip install qdrant-client",
+    )
 
     def __init__(self, package_name, error_message):
         self.package_name = package_name

adalflow/tests/test_qdrant_retriever.py

@@ -0,0 +1,186 @@
+import pytest
+from unittest.mock import MagicMock
+from adalflow.components.retriever import QdrantRetriever
+from adalflow.core.types import (
+    RetrieverOutput,
+    Document,
+)
+from adalflow.core.embedder import Embedder
+
+qdrant_client = pytest.importorskip(
+    "qdrant_client", reason="qdrant_client not installed"
+)
+
+COLLECTION_NAME = "test_collection"
+
+
+@pytest.fixture
+def mock_qdrant_client():
+    return MagicMock(spec=qdrant_client.QdrantClient)
+
+
+@pytest.fixture
+def qdrant_retriever(mock_qdrant_client):
+    return QdrantRetriever(
+        collection_name=COLLECTION_NAME,
+        client=mock_qdrant_client,
+        embedder=MagicMock(spec=Embedder),
+        top_k=5,
+    )
+
+
+def test_reset_index(qdrant_retriever, mock_qdrant_client):
+    mock_qdrant_client.collection_exists.return_value = True
+    qdrant_retriever.reset_index()
+    mock_qdrant_client.delete_collection.assert_called_once_with(COLLECTION_NAME)
+
+
+def test_call_single_query(qdrant_retriever, mock_qdrant_client):
+    query = "test query"
+
+    mock_point = MagicMock()
+    mock_point.id = 1
+    mock_point.score = 0.9
+    mock_point.payload = {"text": "retrieved text", "meta_data": {"key": "value"}}
+    mock_point.vector = [0.1, 0.2, 0.3]
+
+    mock_query_response = MagicMock()
+    mock_query_response.points = [mock_point]
+
+    mock_qdrant_client.query_batch_points.return_value = [mock_query_response]
+
+    result = qdrant_retriever.call(query)
+
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], RetrieverOutput)
+    assert result[0].query == query
+    assert len(result[0].doc_indices) == 1
+    assert result[0].doc_indices[0] == 1
+    assert len(result[0].doc_scores) == 1
+    assert result[0].doc_scores[0] == 0.9
+    assert len(result[0].documents) == 1
+    assert isinstance(result[0].documents[0], Document)
+    assert result[0].documents[0].text == "retrieved text"
+    assert result[0].documents[0].meta_data == {"key": "value"}
+
+
+def test_get_first_vector_name(qdrant_retriever, mock_qdrant_client):
+    # Check single unnamed vector
+    mock_qdrant_client.get_collection.return_value = MagicMock(
+        config=MagicMock(
+            params=MagicMock(
+                vectors=qdrant_client.models.VectorParams(
+                    size=1, distance=qdrant_client.models.Distance.COSINE
+                )
+            )
+        )
+    )
+    vector_name = qdrant_retriever._get_first_vector_name()
+    assert vector_name is None
+
+    mock_qdrant_client.get_collection.return_value = MagicMock(
+        config=MagicMock(
+            params=MagicMock(vectors={"vector1": "details", "vector2": "details"})
+        )
+    )
+    vector_name = qdrant_retriever._get_first_vector_name()
+    assert vector_name == "vector1"
+
+
+def test_points_to_output():
+    # Prepare mocked ScoredPoint
+    mock_point = MagicMock()
+    mock_point.id = 1
+    mock_point.score = 0.9
+    mock_point.payload = {"text": "sample text", "meta_data": {"key": "value"}}
+    mock_point.vector = [0.1, 0.2, 0.3]
+
+    points = [mock_point]
+    query = "test query"
+    text_key = "text"
+    metadata_key = "meta_data"
+    vector_name = "vector_name"
+
+    result = QdrantRetriever._points_to_output(
+        points, query, text_key, metadata_key, vector_name
+    )
+
+    assert isinstance(result, RetrieverOutput)
+    assert result.query == query
+    assert result.doc_indices == [1]
+    assert result.doc_scores == [0.9]
+    assert len(result.documents) == 1
+    assert isinstance(result.documents[0], Document)
+    assert result.documents[0].text == "sample text"
+    assert result.documents[0].meta_data == {"key": "value"}
+    assert result.documents[0].vector == [0.1, 0.2, 0.3]
+
+
+def test_doc_from_point():
+    mock_point = MagicMock()
+    mock_point.id = 1
+    mock_point.payload = {"content": "sample text", "some_meta": {"key": "value"}}
+    mock_point.vector = [0.1, 0.2, 0.3]
+
+    text_key = "content"
+    metadata_key = "some_meta"
+    vector_name = None
+
+    document = QdrantRetriever._doc_from_point(
+        mock_point, text_key, metadata_key, vector_name
+    )
+
+    assert isinstance(document, Document)
+    assert document.id == 1
+    assert document.text == "sample text"
+    assert document.meta_data == {"key": "value"}
+    assert document.vector == [0.1, 0.2, 0.3]
+
+
+def test_doc_from_point_with_vector_name():
+    mock_point = MagicMock()
+    mock_point.id = 1
+    mock_point.payload = {"text": "sample text", "meta_data": {"key": "value"}}
+    mock_point.vector = {"vector_name": [0.4, 0.5, 0.6]}
+
+    text_key = "text"
+    metadata_key = "meta_data"
+    vector_name = "vector_name"
+
+    document = QdrantRetriever._doc_from_point(
+        mock_point, text_key, metadata_key, vector_name
+    )
+
+    assert isinstance(document, Document)
+    assert document.id == 1
+    assert document.text == "sample text"
+    assert document.meta_data == {"key": "value"}
+    assert document.vector == [0.4, 0.5, 0.6]
+
+
+def test_call_with_custom_limit(qdrant_retriever, mock_qdrant_client):
+    query = "test query"
+    custom_limit = 5
+
+    mock_point = MagicMock()
+    mock_point.id = 1
+    mock_point.score = 0.9
+    mock_point.payload = {"text": "retrieved text", "meta_data": {"key": "value"}}
+    mock_point.vector = [0.1, 0.2, 0.3]
+
+    mock_query_response = MagicMock(spec=qdrant_client.models.QueryResponse)
+    mock_query_response.points = [mock_point]
+
+    mock_qdrant_client.query_batch_points.return_value = [mock_query_response]
+
+    qdrant_retriever.call([query, query, query], top_k=custom_limit)
+
+    mock_qdrant_client.query_batch_points.assert_called_once()
+
+    collection_name = mock_qdrant_client.query_batch_points.call_args[0]
+    assert collection_name == (COLLECTION_NAME,)
+
+    requests = mock_qdrant_client.query_batch_points.call_args[1]["requests"]
+    for request in requests:
+        assert request.limit == custom_limit