agentic rag

langchain-streamlit-demo/app.py

@@ -442,6 +442,7 @@ if st.session_state.llm:
                 MEMORY,
                 chat_prompt,
                 prompt,
+                STMEMORY,
             )
 
             # --- LLM call ---

langchain-streamlit-demo/llm_resources.py

@@ -1,8 +1,15 @@
+import uuid
 from tempfile import NamedTemporaryFile
 from typing import Tuple, List, Optional, Dict
 
+from langchain.agents import AgentExecutor
+from langchain.agents.agent_toolkits import create_retriever_tool
+from langchain.agents.openai_functions_agent.agent_token_buffer_memory import (
+    AgentTokenBufferMemory,
+)
+from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.chains import RetrievalQA, LLMChain
+from langchain.chains import LLMChain
 from langchain.chat_models import (
     AzureChatOpenAI,
     ChatOpenAI,
@@ -11,15 +18,16 @@ from langchain.chat_models import (
 )
 from langchain.document_loaders import PyPDFLoader
 from langchain.embeddings import AzureOpenAIEmbeddings, OpenAIEmbeddings
+from langchain.prompts import MessagesPlaceholder
 from langchain.retrievers import EnsembleRetriever
-from langchain.schema import Document, BaseRetriever
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-
 from langchain.retrievers.multi_query import MultiQueryRetriever
 from langchain.retrievers.multi_vector import MultiVectorRetriever
+from langchain.schema import Document, BaseRetriever
+from langchain.schema.runnable import RunnablePassthrough
 from langchain.storage import InMemoryStore
-import uuid
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import FAISS
+from langchain_core.messages import SystemMessage
 
 from defaults import DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP, DEFAULT_RETRIEVER_K
 from qagen import get_rag_qa_gen_chain
@@ -34,6 +42,7 @@ def get_runnable(
     memory,
     chat_prompt,
     summarization_prompt,
+    chat_history,
 ):
     if not use_document_chat:
         return LLMChain(
@@ -54,13 +63,44 @@ def get_runnable(
             llm,
         )
     else:
-        return RetrievalQA.from_chain_type(
+        tool = create_retriever_tool(
+            retriever,
+            "search_user_document",
+            "Retrieves custom context provided by the user for this conversation. Use this if you cannot answer immediately and confidently.",
+        )
+        tools = [tool]
+        memory_key = "agent_history"
+        system_message = SystemMessage(
+            content=(
+                "Do your best to answer the questions. "
+                "Feel free to use any tools available to look up "
+                "relevant information, only if necessary"
+            ),
+        )
+        prompt = OpenAIFunctionsAgent.create_prompt(
+            system_message=system_message,
+            extra_prompt_messages=[MessagesPlaceholder(variable_name=memory_key)],
+        )
+        agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
+
+        agent_memory = AgentTokenBufferMemory(
+            chat_memory=chat_history,
+            memory_key=memory_key,
             llm=llm,
-            chain_type=document_chat_chain_type,
-            retriever=retriever,
-            memory=memory,
-            output_key="output_text",
-        ) | (lambda output: output["output_text"])
+        )
+
+        agent_executor = AgentExecutor(
+            agent=agent,
+            tools=tools,
+            memory=agent_memory,
+            verbose=True,
+            return_intermediate_steps=True,
+        )
+        return (
+            {"input": RunnablePassthrough()}
+            | agent_executor
+            | (lambda output: output["output"])
+        )
 
 
 def get_llm(