support langchain new version and set agent False as default (#90)

Signed-off-by: ChengZi <[email protected]>

main.py

@@ -3,6 +3,7 @@
 import uuid
 
 import uvicorn
+from functools import partial
 from fastapi import FastAPI, UploadFile
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import PlainTextResponse
@@ -17,6 +18,9 @@
 parser.add_argument('--langchain', action='store_true')
 parser.add_argument('--towhee', action='store_true')
 parser.add_argument('--moniter', action='store_true')
+parser.add_argument('--agent', action='store_true',
+ help='The default is False, which only works when `--langchain` is enabled.'
+ ' It means using the agent in langchain to dynamically select tools.')
 parser.add_argument('--max_observation', default=1000)
 parser.add_argument('--name', default=str(uuid.uuid4()))
 args = parser.parse_args()
@@ -29,13 +33,15 @@
 USE_TOWHEE = args.towhee
 MAX_OBSERVATION = args.max_observation
 ENABLE_MONITER = args.moniter
+ENABLE_AGENT = args.agent
 NAME = args.name
 
 assert (USE_LANGCHAIN and not USE_TOWHEE ) or (USE_TOWHEE and not USE_LANGCHAIN), \
  'The service should start with either "--langchain" or "--towhee".'
 
 if USE_LANGCHAIN:
  from src_langchain.operations import chat, insert, drop, check, get_history, clear_history, count # pylint: disable=C0413
+ chat = partial(chat, enable_agent=ENABLE_AGENT)
 if USE_TOWHEE:
  from src_towhee.operations import chat, insert, drop, check, get_history, clear_history, count # pylint: disable=C0413
 if ENABLE_MONITER:

requirements.txt

@@ -1,4 +1,4 @@
-langchain==0.0.230
+langchain==0.0.322
 unstructured
 pexpect
 pdf2image

src_langchain/agent/chat_agent.py

@@ -1,11 +1,11 @@
 from typing import Any, List, Optional, Sequence
-from pydantic import Field
 
-from langchain.agents.conversational_chat.prompt import PREFIX, SUFFIX
+from pydantic import Field
+from langchain.schema.language_model import BaseLanguageModel
 from langchain.callbacks.base import BaseCallbackManager
 from langchain.agents import ConversationalChatAgent, AgentOutputParser, Agent
 from langchain.prompts.base import BasePromptTemplate
-from langchain.schema import BaseOutputParser, BaseLanguageModel
+from langchain.schema import BaseOutputParser
 from langchain.tools.base import BaseTool
 
 from .output_parser import OutputParser

src_langchain/data_loader/data_parser.py

@@ -42,7 +42,7 @@ def __call__(self, data_src, source_type: str = 'file') -> List[str]:
  token_count += len(self.enc.encode(doc))
  return docs, token_count
 
- def from_files(self, files: list, encoding: Optional[str] = None) -> List[Document]:
+ def from_files(self, files: list, encoding: Optional[str] = 'utf-8') -> List[Document]:
  '''Load documents from path or file-like object, return a list of unsplit LangChain Documents'''
  docs = []
  for file in files:

src_langchain/operations.py

@@ -4,6 +4,7 @@
 from typing import List
 
 from langchain.agents import Tool, AgentExecutor
+from langchain.chains import ConversationalRetrievalChain
 
 sys.path.append(os.path.dirname(__file__))
 
@@ -13,41 +14,52 @@
 from store import MemoryStore, DocStore # pylint: disable=C0413
 from data_loader import DataParser # pylint: disable=C0413
 
-
 logger = logging.getLogger(__name__)
 
 encoder = TextEncoder()
 chat_llm = ChatLLM()
 load_data = DataParser()
 
 
-def chat(session_id, project, question):
+def chat(session_id, project, question, enable_agent=False):
  '''Chat API'''
  doc_db = DocStore(
  table_name=project,
  embedding_func=encoder,
  )
  memory_db = MemoryStore(table_name=project, session_id=session_id)
 
- tools = [
- Tool(
- name='Search',
- func=doc_db.search,
- description='Search through Milvus.'
+ if enable_agent: # use agent
+ memory_db.memory.output_key = None
+ tools = [
+ Tool(
+ name='Search',
+ func=doc_db.search,
+ description='useful for search professional knowledge and information'
+ )
+ ]
+ agent = ChatAgent.from_llm_and_tools(llm=chat_llm, tools=tools)
+ agent_chain = AgentExecutor.from_agent_and_tools(
+ agent=agent,
+ tools=tools,
+ memory=memory_db.memory,
+ verbose=False
  )
- ]
- agent = ChatAgent.from_llm_and_tools(llm=chat_llm, tools=tools)
- agent_chain = AgentExecutor.from_agent_and_tools(
- agent=agent,
- tools=tools,
- memory=memory_db.memory,
- verbose=False
- )
- try:
- final_answer = agent_chain.run(input=question)
- return question, final_answer
- except Exception as e: # pylint: disable=W0703
- return question, f'Something went wrong:\n{e}'
+ try:
+ final_answer = agent_chain.run(input=question)
+ return question, final_answer
+ except Exception as e: # pylint: disable=W0703
+ return question, f'Something went wrong:\n{e}'
+ else: # use chain
+ memory_db.memory.output_key = 'answer'
+ qa = ConversationalRetrievalChain.from_llm(
+ llm=chat_llm,
+ retriever=doc_db.vector_db.as_retriever(),
+ memory=memory_db.memory,
+ return_generated_question=True
+ )
+ qa_result = qa(question)
+ return qa_result['generated_question'], qa_result['answer']
 
 
 def insert(data_src, project, source_type: str = 'file'):
@@ -93,6 +105,7 @@ def check(project):
  raise RuntimeError from e
  return {'store': doc_check, 'memory': memory_check}
 
+
 def count(project):
  '''Count entities.'''
  try:
@@ -130,7 +143,6 @@ def load(document_strs: List[str], project: str):
  num = doc_db.insert(document_strs)
  return num
 
-
 # if __name__ == '__main__':
 # project = 'akcio'
 # data_src = 'https://docs.towhee.io/'