Fix & Improve (see details)

1. The agent llm were used with `max_tokens=256`(default), we convert
   llm be gpt-3.5-turbo-16k with `max_tokens=6000`

2. After expand the max_tokens, the agent response(thought) will not incomplete.

Author: Lin-jun-xiang

agent/agent.py

@@ -6,7 +6,7 @@
 from langchain.agents import AgentType, Tool, initialize_agent
 from langchain.callbacks import get_openai_callback
 from langchain.chains import LLMChain
-from langchain.llms import OpenAI
+from langchain.chat_models import ChatOpenAI
 from langchain.prompts import PromptTemplate
 
 openai.api_key = os.getenv('OPENAI_API_KEY')
@@ -16,7 +16,11 @@
 class AgentHelper:
     """Add agent to help docGPT can be perfonm better."""
     def __init__(self) -> None:
-        self.llm = OpenAI(temperature=0)
+        self.llm = ChatOpenAI(
+            temperature=0.2,
+            max_tokens=6000,
+            model_name='gpt-3.5-turbo-16k'
+        )
         self.agent_ = None
         self.tools = []
 
@@ -32,7 +36,6 @@ def get_calculate_chain(self) -> Tool:
 
     @property
     def get_searp_chain(self) -> Tool:
-
         search = SerpAPIWrapper()
         tool = Tool(
             name='Search',
@@ -47,7 +50,7 @@ def create_doc_chat(self, docGPT) -> Tool:
             name='DocumentGPT',
             func=docGPT.run,
             description="""
-            useful for when you need to answer questions from the context of PDF,
+            useful for when you need to answer questions from the context of PDF
             """
         )
         return tool
@@ -58,12 +61,12 @@ def create_llm_chain(self) -> Tool:
             input_variables = ['query'],
             template = '{query}'
         )
-        llm_chain = LLMChain(llm=self.llm, prompt = prompt)
+        llm_chain = LLMChain(llm=self.llm, prompt=prompt)
 
         tool = Tool(
             name='LLM',
             func=llm_chain.run,
-            description='useful for general purpose queries and logic'
+            description='useful for general purpose queries and logic.'
         )
         return tool
 

app.py

@@ -14,7 +14,6 @@
 from docGPT import DocGPT
 from model import PDFLoader
 
-
 langchain.llm_cache = InMemoryCache()
 
 OPENAI_API_KEY = ''
@@ -82,7 +81,7 @@ def load_api_key() -> None:
         temp_file_path = temp_file.name
 
         docs = PDFLoader.load_documents(temp_file_path)
-        docs = PDFLoader.split_documents(docs, chunk_size=2500, chunk_overlap=200)
+        docs = PDFLoader.split_documents(docs, chunk_size=2000, chunk_overlap=200)
 
         temp_file.close()
         if temp_file_path:
@@ -113,7 +112,8 @@ def load_api_key() -> None:
             tools = [
                 docGPT_tool,
                 search_tool,
-                llm_tool
+                # llm_tool, # This will cause agent confuse
+                calculate_tool
             ]
             agent_.initialize(tools)
         except Exception as e:

docGPT/docGPT.py

@@ -4,13 +4,11 @@
 import openai
 from langchain.callbacks import get_openai_callback
 from langchain.chains import ConversationalRetrievalChain, RetrievalQA
+from langchain.chat_models import ChatOpenAI
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.llms import OpenAI
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
 from langchain.vectorstores import Chroma
-from langchain.chat_models import ChatOpenAI
-
 
 openai.api_key = os.getenv('OPENAI_API_KEY')
 
@@ -84,8 +82,8 @@ def __init__(self, docs):
         self.qa_chain = None
         self.llm = ChatOpenAI(
             temperature=0.2,
-            max_tokens=2000,
-            model_name='gpt-3.5-turbo'
+            max_tokens=6000,
+            model_name='gpt-3.5-turbo-16k'
         )
 
         self.prompt_template = """

model/data_connection.py

@@ -1,21 +1,21 @@
 import json
 import os
+from typing import Iterator
 
 from langchain.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 
 class PDFLoader:
     @staticmethod
-    def get_pdf_files(path: str) -> list:
-        if path.endswith('.pdf'):
-            return f'./PDF/uploaded/{path}'
-    
-        else:
-                file_names = os.listdir(f'./PDF/{path}')
-                pdf_files = [name for name in file_names if name.endswith('.pdf')]
-
-        return pdf_files
+    def get_pdf_files(path: str) -> Iterator[str]:
+        try:
+            yield from [
+                file_name for file_name in os.listdir(f'{path}')
+                if file_name.endswith('.pdf')
+            ]
+        except FileNotFoundError as e:
+            print(f'\033[31m{e}')
 
     @staticmethod
     def load_documents(pdf_file: str) -> PyMuPDFLoader: