From 1a46ea481666fce52e2b28d0ff103016a1b09e57 Mon Sep 17 00:00:00 2001 From: Romain Quinet Date: Sat, 7 Oct 2023 08:48:08 +0200 Subject: [PATCH] Improved chat mode --- main.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/main.py b/main.py index 4368820..41f62d6 100644 --- a/main.py +++ b/main.py @@ -4,9 +4,13 @@ from llama_index.node_parser import SimpleNodeParser from llama_index import VectorStoreIndex from llama_index.llms import OpenAI, ChatMessage, MessageRole from llama_index.prompts import ChatPromptTemplate +from llama_index import set_global_handler +from llama_index.chat_engine.types import ChatMode import os import re +#set_global_handler("simple") + llm = OpenAI(model="gpt-4", temperature=0, max_tokens=256) service_context = ServiceContext.from_defaults(llm=llm) set_global_service_context(service_context) @@ -90,16 +94,13 @@ refine_template = ChatPromptTemplate(chat_refine_msgs) chat_engine = index.as_chat_engine( text_qa_template=text_qa_template, - refine_template=refine_template + refine_template=refine_template, + chat_mode=ChatMode.OPENAI ) while True: try: - user_prompt = input("Prompt: ") - streaming_response = chat_engine.stream_chat(user_prompt) - for token in streaming_response.response_gen: - print(token, end="") - print("\n") + chat_engine.chat_repl() except KeyboardInterrupt: break