feat: add StreamingResponse support for ConversationalRetrievalChain chain type

examples/README.md

@@ -34,6 +34,7 @@ To run a demo example, select the command based on the langchain use case you wa
 
 - Conversation Chain: `uvicorn app.conversation_chain:app --reload`
 - Retrieval QA with Sources Chain: `uvicorn app.retrieval_qa_chain:app --reload`
+- Conversational Retrieval: `uvicorn app.conversational_retrieval:app --reload`
 
 You can also use the "Run & Debug" VSCode feature to run one of the applications.
 

examples/app/conversational_retrieval.py

@@ -0,0 +1,84 @@
+from functools import lru_cache
+from typing import Callable
+
+from dotenv import load_dotenv
+from fastapi import Depends, FastAPI
+from fastapi.templating import Jinja2Templates
+from langchain.chains import ConversationalRetrievalChain, LLMChain
+from langchain.chains.question_answering import load_qa_chain
+from langchain.chat_models import ChatOpenAI
+from pydantic import BaseModel
+
+from fastapi_async_langchain.responses import ConversationalRetrievalStreamingResponse
+from fastapi_async_langchain.testing import mount_gradio_app
+
+load_dotenv()
+
+app = mount_gradio_app(FastAPI(title="ConversationalRetrievalChainDemo"))
+
+templates = Jinja2Templates(directory="templates")
+
+
+class QueryRequest(BaseModel):
+    query: str
+    history: list[list[str]] = []
+
+
+def conversational_retrieval_chain_dependency() -> (
+    Callable[[], ConversationalRetrievalChain]
+):
+    @lru_cache(maxsize=1)
+    def dependency() -> ConversationalRetrievalChain:
+        from langchain.chains.conversational_retrieval.prompts import (
+            CONDENSE_QUESTION_PROMPT,
+        )
+        from langchain.embeddings import OpenAIEmbeddings
+        from langchain.vectorstores import FAISS
+
+        db = FAISS.load_local(
+            folder_path="vector_stores/",
+            index_name="langchain-python",
+            embeddings=OpenAIEmbeddings(),
+        )
+
+        question_generator = LLMChain(
+            llm=ChatOpenAI(
+                temperature=0,
+                streaming=True,
+            ),
+            prompt=CONDENSE_QUESTION_PROMPT,
+        )
+        doc_chain = load_qa_chain(
+            llm=ChatOpenAI(
+                temperature=0,
+                streaming=True,
+            ),
+            chain_type="stuff",
+        )
+
+        return ConversationalRetrievalChain(
+            combine_docs_chain=doc_chain,
+            question_generator=question_generator,
+            retriever=db.as_retriever(),
+            return_source_documents=True,
+            verbose=True,
+        )
+
+    return dependency
+
+
+conversational_retrieval_chain = conversational_retrieval_chain_dependency()
+
+
+@app.post("/chat")
+async def chat(
+    request: QueryRequest,
+    chain: ConversationalRetrievalChain = Depends(conversational_retrieval_chain),
+) -> ConversationalRetrievalStreamingResponse:
+    inputs = {
+        "question": request.query,
+        "chat_history": [(human, ai) for human, ai in request.history],
+    }
+    return ConversationalRetrievalStreamingResponse.from_chain(
+        chain, inputs, media_type="text/event-stream"
+    )

fastapi_async_langchain/responses/__init__.py

@@ -1,4 +1,9 @@
+from .conversational_retrieval import ConversationalRetrievalStreamingResponse
 from .llm import LLMChainStreamingResponse
 from .retrieval_qa import RetrievalQAStreamingResponse
 
-__all__ = ["LLMChainStreamingResponse", "RetrievalQAStreamingResponse"]
+__all__ = [
+    "LLMChainStreamingResponse",
+    "RetrievalQAStreamingResponse",
+    "ConversationalRetrievalStreamingResponse",
+]

fastapi_async_langchain/responses/conversational_retrieval.py

@@ -0,0 +1,22 @@
+from typing import Any, Awaitable, Callable, Dict, Union
+
+from langchain.chains.retrieval_qa.base import BaseRetrievalQA
+from starlette.types import Send
+
+from ..callbacks import AsyncRetrievalQAStreamingCallback
+from .base import BaseLangchainStreamingResponse
+
+
+class ConversationalRetrievalStreamingResponse(BaseLangchainStreamingResponse):
+    """BaseLangchainStreamingResponse class wrapper for ConversationalRetrievalStreamingResponse instances."""
+
+    @staticmethod
+    def _create_chain_executor(
+        chain: BaseRetrievalQA, inputs: Union[Dict[str, Any], Any]
+    ) -> Callable[[Send], Awaitable[Any]]:
+        async def wrapper(send: Send):
+            return await chain.acall(
+                inputs=inputs, callbacks=[AsyncRetrievalQAStreamingCallback(send=send)]
+            )
+
+        return wrapper

fastapi_async_langchain/websockets/__init__.py

@@ -1,9 +1,11 @@
 from .base import BaseLangchainWebsocketConnection
+from .conversational_retrieval import ConversationalRetrievalWebsocketConnection
 from .llm import LLMChainWebsocketConnection
 from .retrieval_qa import RetrievalQAWebsocketConnection
 
 __all__ = [
     "BaseLangchainWebsocketConnection",
     "LLMChainWebsocketConnection",
     "RetrievalQAWebsocketConnection",
+    "ConversationalRetrievalWebsocketConnection",
 ]

pyproject.toml

@@ -9,7 +9,9 @@ packages = [{include = "fastapi_async_langchain"}]
 [tool.poetry.dependencies]
 python = "^3.9"
 fastapi = "^0.95.1"
-langchain = "^0.0.157"
+langchain = "^0.0.164"
+urllib3 = "<=1.26.15"  # added due to poetry errors
+python-dotenv = "^1.0.0"
 
 [build-system]
 requires = ["poetry-core"]