Merge pull request 'Change RetGenLangGraph to use streaming instead of invoking on the LLM' (#17) from reg_gen_stream into main

Reviewed-on: AI_team/generic-RAG-demo#17

generic_rag/app.py

@@ -98,9 +98,8 @@ async def process_response(message):
 
     chainlit_response = cl.Message(content="")
 
-    response = graph.invoke(message.content, config=config)
-
-    await chainlit_response.stream_token(f"{response}\n")
+    async for response in graph.stream(message.content, config=config):
+        await chainlit_response.stream_token(response)
 
     pdf_sources = graph.get_last_pdf_sources()
     if len(pdf_sources) > 0:

generic_rag/graphs/ret_gen.py

@@ -1,9 +1,12 @@
 import logging
 from pathlib import Path
-from typing import Any, Union
+from typing import Any, AsyncGenerator
 
 from langchain import hub
+from langchain_chroma import Chroma
 from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, START, StateGraph
 from typing_extensions import List, TypedDict
@@ -19,7 +22,7 @@ class State(TypedDict):
 
 
 class RetGenLangGraph:
-    def __init__(self, vector_store, chat_model, embedding_model):
+    def __init__(self, vector_store: Chroma, chat_model: BaseChatModel, embedding_model: Embeddings):
         self.vector_store = vector_store
         self.chat_model = chat_model
         self.embedding_model = embedding_model
@@ -32,21 +35,21 @@ class RetGenLangGraph:
         graph_builder.add_edge("_generate", END)
 
         self.graph = graph_builder.compile(memory)
-        self.last_invoke = None
+        self.last_retrieved_docs = []
 
-    def invoke(self, message: str, config: dict) -> Union[dict[str, Any], Any]:
-        self.last_invoke = self.graph.invoke({"question": message}, config=config)
-        return self.last_invoke["answer"]
+    async def stream(self, message: str, config: dict) -> AsyncGenerator[Any, Any]:
+        async for response, _ in self.graph.astream({"question": message}, stream_mode="messages", config=config):
+            yield response.content
 
     def _retrieve(self, state: State) -> dict:
-        retrieved_docs = self.vector_store.similarity_search(state["question"])
-        return {"context": retrieved_docs}
+        self.last_retrieved_docs = self.vector_store.similarity_search(state["question"])
+        return {"context": self.last_retrieved_docs}
 
-    def _generate(self, state: State) -> dict:
+    async def _generate(self, state: State) -> AsyncGenerator[Any, Any]:
         docs_content = "\n\n".join(doc.page_content for doc in state["context"])
-        messages = self.prompt.invoke({"question": state["question"], "context": docs_content})
-        response = self.chat_model.invoke(messages)
-        return {"answer": response.content}
+        messages = await self.prompt.ainvoke({"question": state["question"], "context": docs_content})
+        async for response in self.chat_model.astream(messages):
+            yield {"answer": response.content}
 
     def get_last_pdf_sources(self) -> dict[str, list[int]]:
         """
@@ -54,30 +57,30 @@ class RetGenLangGraph:
         """
         pdf_sources = {}
 
-        if self.last_invoke is None:
+        if not self.last_retrieved_docs:
             return pdf_sources
 
-        for context in self.last_invoke["context"]:
+        for doc in self.last_retrieved_docs:
             try:
-                Path(context.metadata["source"]).suffix == ".pdf"
+                Path(doc.metadata["source"]).suffix == ".pdf"
             except KeyError:
                 continue
             else:
-                source = context.metadata["source"]
+                source = doc.metadata["source"]
 
             if source not in pdf_sources:
                 pdf_sources[source] = set()
 
             # The page numbers are in the `page_numer` and `page` fields.
             try:
-                page_number = context.metadata["page_number"]
+                page_number = doc.metadata["page_number"]
             except KeyError:
                 pass
             else:
                 pdf_sources[source].add(page_number)
 
             try:
-                page_number = context.metadata["page"]
+                page_number = doc.metadata["page"]
             except KeyError:
                 pass
             else:
@@ -94,15 +97,14 @@ class RetGenLangGraph:
         """
         web_sources = set()
 
-        if self.last_invoke is None:
+        if not self.last_retrieved_docs:
             return web_sources
 
-        for context in self.last_invoke["context"]:
+        for doc in self.last_retrieved_docs:
             try:
-                context.metadata["filetype"] == "web"
+                if doc.metadata["filetype"] == "web":
+                    web_sources.add(doc.metadata["source"])
             except KeyError:
                 continue
-            else:
-                web_sources.add(context.metadata["source"])
 
         return web_sources