improved main agent/doc agent interaction

AI_CHANGELOG.md

@@ -1,4 +1,10 @@
 # AI CHANGELOG
+## [Refactor Memory Management and Remove Intermediate Steps](https://github.com/joshuasundance-swca/langchain-streamlit-demo/commit/3a2b820bcfb7ef8dc4157626de454be9267e1e91)
+Tue Dec 19 09:35:23 2023 -0500
+- This commit refactors the memory management in the Streamlit app and other resources. It replaces the 'AgentTokenBufferMemory' with 'ConversationBufferMemory' for a more efficient memory usage. The 'return_messages' flag is set to 'True' in the new memory management system.
+- The commit also disables the return of intermediate steps in both the Streamlit app and agent execution by setting 'return_intermediate_steps' to 'False'. This change can help to improve performance by reducing unnecessary computations.
+- Furthermore, the commit simplifies the 'python_coder.py' script by removing unnecessary imports and argument parsing related to 'black', 'ruff', and 'mypy'. It also removes the main execution block, indicating that this script may now be used as a module rather than a standalone script.
+- The changes in this commit aim to improve performance and maintainability of the code.
 ## [Updated application version to 2.1.2](https://github.com/joshuasundance-swca/langchain-streamlit-demo/commit/3fddfc9d34f66f92255f6a0d0878e20db401d694)
 Mon Dec 18 11:02:13 2023 -0500
 - Bumped the current version of the application in the bumpver.toml file from 2.1.1 to 2.1.2.

langchain-streamlit-demo/app.py

@@ -431,7 +431,7 @@ if st.session_state.llm:
                     callbacks=callbacks,
                     tags=["Streamlit Chat"],
                     verbose=True,
-                    return_intermediate_steps=True,
+                    return_intermediate_steps=False,
                 )
                 if st.session_state.provider == "Anthropic":
                     config["max_concurrency"] = 5
@@ -507,17 +507,17 @@ if st.session_state.llm:
                     doc_chain_agent = get_doc_agent(
                         [doc_chain_tool],
                     )
-                    doc_agent_tool = Tool.from_function(
+                    doc_question_tool = Tool.from_function(
                         func=lambda s: doc_chain_agent.invoke(
                             s,
                         ),
-                        name="document-agent",
-                        description="this assistant returns a response based on the user's custom context. "
+                        name="document-question-tool",
+                        description="this assistant answers a question based on the user's custom context. "
+                        "this assistant responds to fully formed questions."
                         "if the user's meaning is unclear, perhaps the answer is here. "
-                        "generally speaking, try this tool before conducting web research."
-                        "it is best to send this tool a question, as it will attempt to break complex questions down into several, simpler questions.",
+                        "generally speaking, try this tool before conducting web research.",
                     )
-                    TOOLS = [doc_agent_tool, research_assistant_tool] + default_tools
+                    TOOLS = [doc_question_tool, research_assistant_tool] + default_tools
 
                 st.session_state.chain = get_agent(
                     TOOLS,

langchain-streamlit-demo/llm_resources.py

@@ -3,9 +3,6 @@ from tempfile import NamedTemporaryFile
 from typing import Tuple, List, Optional, Dict
 
 from langchain.agents import AgentExecutor, AgentType, initialize_agent
-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 LLMChain
@@ -57,10 +54,17 @@ def get_agent(
     )
     agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
 
-    agent_memory = AgentTokenBufferMemory(
+    # agent_memory = AgentTokenBufferMemory(
+    #     chat_memory=chat_history,
+    #     memory_key=memory_key,
+    #     llm=llm,
+    # )
+    from langchain.memory import ConversationBufferMemory
+
+    agent_memory = ConversationBufferMemory(
         chat_memory=chat_history,
+        return_messages=True,
         memory_key=memory_key,
-        llm=llm,
     )
 
     agent_executor = AgentExecutor(
@@ -68,7 +72,7 @@ def get_agent(
         tools=tools,
         memory=agent_memory,
         verbose=True,
-        return_intermediate_steps=True,
+        return_intermediate_steps=False,
         callbacks=callbacks,
     )
     return (

langchain-streamlit-demo/python_coder.py

@@ -1,16 +1,14 @@
 """langchain python coder-- requires black, ruff, and mypy."""
 
-import argparse
 import os
 import re
 import subprocess  # nosec
 import tempfile
-from importlib.util import find_spec
 
 from langchain.agents import initialize_agent, AgentType
 from langchain.agents.tools import Tool
-from langchain.chat_models import ChatOpenAI
 from langchain.llms.base import BaseLLM
+from langchain.memory import ConversationBufferMemory
 from langchain.prompts import ChatPromptTemplate
 from langchain.pydantic_v1 import BaseModel, validator, Field, ValidationError
 
@@ -177,35 +175,17 @@ def get_agent(
     llm: BaseLLM,
     agent_type: AgentType = AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
 ):
+    agent_memory = ConversationBufferMemory(
+        return_messages=True,
+        memory_key="chat_history",
+    )
     return initialize_agent(
         tools,
         llm,
         agent=agent_type,
         verbose=True,
         handle_parsing_errors=True,
+        memory=agent_memory,
         prompt=prompt,
-        # return_intermediate_steps=True,
+        return_intermediate_steps=False,
     ) | (lambda output: output["output"])
-
-
-if __name__ == "__main__":
-    for lib in ("black", "ruff", "mypy"):
-        if find_spec(lib) is None:
-            raise ImportError(
-                "You must install black, ruff, and mypy to use this tool. "
-                "You can install them with `pip install black ruff mypy`.",
-            )
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", "-m", default="gpt-4-1106-preview")
-    parser.add_argument("instruction")
-    args = parser.parse_args()
-
-    instruction = args.instruction
-    model = args.model
-
-    llm = ChatOpenAI(model_name=model, temperature=0.0)
-    agent = get_agent(llm)
-
-    output = agent.invoke({"input": instruction}, config=dict(verbose=True))
-    print(output)