From 3295bb8992bdca8c11f223ddd31a97fd57e3fa60 Mon Sep 17 00:00:00 2001
From: Ruben Lucas <ruben.lucas@sogeti.com>
Date: Wed, 9 Apr 2025 15:23:54 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Find=20and=20add=20found=20sources?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 generic_rag/app.py                 | 15 +++++++++++
 generic_rag/graphs/cond_ret_gen.py | 42 ++++++++++++++++++++++--------
 2 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/generic_rag/app.py b/generic_rag/app.py
index 56b7162..64ec26d 100644
--- a/generic_rag/app.py
+++ b/generic_rag/app.py
@@ -129,6 +129,21 @@ async def process_cond_response(message):
     for response in graph.stream(message.content, config=config):
         await chainlit_response.stream_token(response)
 
+    if len(graph.last_retrieved_docs) > 0:
+        await chainlit_response.stream_token("\nThe following PDF source were consulted:\n")
+        for source, page_numbers in graph.last_retrieved_docs.items():
+            page_numbers = list(page_numbers)
+            page_numbers.sort()
+            # display="side" seems to be not supported by chainlit for PDF's, so we use "inline" instead.
+            chainlit_response.elements.append(cl.Pdf(name="pdf", display="inline", path=source, page=page_numbers[0]))
+            await chainlit_response.update()
+            await chainlit_response.stream_token(f"- '{source}' on page(s): {page_numbers}\n")
+
+    if len(graph.last_retrieved_sources) > 0:
+        await chainlit_response.stream_token("\nThe following web sources were consulted:\n")
+        for source in graph.last_retrieved_sources:
+            await chainlit_response.stream_token(f"- {source}\n")
+
     await chainlit_response.send()
 
 
diff --git a/generic_rag/graphs/cond_ret_gen.py b/generic_rag/graphs/cond_ret_gen.py
index fbbf534..9b0cdd3 100644
--- a/generic_rag/graphs/cond_ret_gen.py
+++ b/generic_rag/graphs/cond_ret_gen.py
@@ -1,5 +1,8 @@
 import logging
-from typing import Any, Iterator, List
+from typing import Any, Iterator
+import re
+import ast
+from pathlib import Path
 
 from langchain_chroma import Chroma
 from langchain_core.documents import Document
@@ -7,6 +10,7 @@ from langchain_core.embeddings import Embeddings
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
 from langchain_core.tools import tool
+from langchain_core.runnables.config import RunnableConfig
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import END, MessagesState, StateGraph
 from langgraph.prebuilt import InjectedStore, ToolNode, tools_condition
@@ -39,28 +43,44 @@ class CondRetGenLangGraph:
 
         self.graph = graph_builder.compile(checkpointer=memory, store=vector_store)
 
-    def stream(self, message: str, config=None) -> Iterator[str]:
+        self.file_path_pattern = r"'file_path'\s*:\s*'((?:[^'\\]|\\.)*)'"
+        self.source_pattern = r"'source'\s*:\s*'((?:[^'\\]|\\.)*)'"
+        self.page_pattern = r"'page'\s*:\s*(\d+)"
+        self.pattern = r"Source:\s*(\{.*?\})"
+
+        self.last_retrieved_docs = {}
+        self.last_retrieved_sources = set()
+
+    def stream(self, message: str, config: RunnableConfig | None = None) -> Iterator[str]:
         for llm_response, metadata in self.graph.stream(
             {"messages": [{"role": "user", "content": message}]}, stream_mode="messages", config=config
         ):
-            if (
-                llm_response.content
-                and not isinstance(llm_response, HumanMessage)
-                and metadata["langgraph_node"] == "_generate"
-            ):
+            if llm_response.content and metadata["langgraph_node"] == "_generate":
                 yield llm_response.content
-
-            # TODO: read souces used in AIMessages and set internal value sources used in last received stream.
+            elif llm_response.name == "_retrieve":
+                dictionary_strings = re.findall(
+                    self.pattern, llm_response.content, re.DOTALL
+                )  # Use re.DOTALL if dicts might span newlines
+                for dict_str in dictionary_strings:
+                    parsed_dict = ast.literal_eval(dict_str)
+                    print(parsed_dict)
+                    if "filetype" in parsed_dict and parsed_dict["filetype"] == "web":
+                        self.last_retrieved_sources.add(parsed_dict["source"])
+                    elif Path(parsed_dict["source"]).suffix == ".pdf":
+                        if parsed_dict["source"] in self.last_retrieved_docs:
+                            self.last_retrieved_docs[parsed_dict["source"]].add(parsed_dict["page"])
+                        else:
+                            self.last_retrieved_docs[parsed_dict["source"]] = {parsed_dict["page"]}
 
     @tool(response_format="content_and_artifact")
     def _retrieve(
         query: str, full_user_content: str, vector_store: Annotated[Any, InjectedStore()]
-    ) -> tuple[str, List[Document]]:
+    ) -> tuple[str, list[Document]]:
         """
         Retrieve information related to a query and user content.
         """
         # This method is used as a tool in the graph.
-        # It's doc-string is used for the pydentic model, please consider doc-string text carefully.
+        # It's doc-string is used for the pydantic model, please consider doc-string text carefully.
         # Furthermore, it can not and should not have the `self` parameter.
         # If you want to pass on state, please refer to:
         # https://python.langchain.com/docs/concepts/tools/#special-type-annotations