Uses a retriever to retrieve a context, set the context in the system prompt,
and then uses an LLM to generate a response, for a fluid chat experience.
Source code in llama-index-core/llama_index/core/chat_engine/context.py
classContextChatEngine(BaseChatEngine):""" Context Chat Engine. Uses a retriever to retrieve a context, set the context in the system prompt, and then uses an LLM to generate a response, for a fluid chat experience. """def__init__(self,retriever:BaseRetriever,llm:LLM,memory:BaseMemory,prefix_messages:List[ChatMessage],node_postprocessors:Optional[List[BaseNodePostprocessor]]=None,context_template:Optional[Union[str,PromptTemplate]]=None,context_refine_template:Optional[Union[str,PromptTemplate]]=None,callback_manager:Optional[CallbackManager]=None,)->None:self._retriever=retrieverself._llm=llmself._memory=memoryself._prefix_messages=prefix_messagesself._node_postprocessors=node_postprocessorsor[]context_template=context_templateorDEFAULT_CONTEXT_TEMPLATEifisinstance(context_template,str):context_template=PromptTemplate(context_template)self._context_template=context_templatecontext_refine_template=context_refine_templateorDEFAULT_REFINE_TEMPLATEifisinstance(context_refine_template,str):context_refine_template=PromptTemplate(context_refine_template)self._context_refine_template=context_refine_templateself.callback_manager=callback_managerorCallbackManager([])fornode_postprocessorinself._node_postprocessors:node_postprocessor.callback_manager=self.callback_manager@classmethoddeffrom_defaults(cls,retriever:BaseRetriever,chat_history:Optional[List[ChatMessage]]=None,memory:Optional[BaseMemory]=None,system_prompt:Optional[str]=None,prefix_messages:Optional[List[ChatMessage]]=None,node_postprocessors:Optional[List[BaseNodePostprocessor]]=None,context_template:Optional[Union[str,PromptTemplate]]=None,context_refine_template:Optional[Union[str,PromptTemplate]]=None,llm:Optional[LLM]=None,**kwargs:Any,)->"ContextChatEngine":"""Initialize a ContextChatEngine from default parameters."""llm=llmorSettings.llmchat_history=chat_historyor[]memory=memoryorChatMemoryBuffer.from_defaults(chat_history=chat_history,token_limit=llm.metadata.context_window-256)ifsystem_promptisnotNone:ifprefix_messagesisnotNone:raiseValueError("Cannot specify both system_prompt and prefix_messages")prefix_messages=[ChatMessage(content=system_prompt,role=llm.metadata.system_role)]prefix_messages=prefix_messagesor[]node_postprocessors=node_postprocessorsor[]returncls(retriever,llm=llm,memory=memory,prefix_messages=prefix_messages,node_postprocessors=node_postprocessors,callback_manager=Settings.callback_manager,context_template=context_template,context_refine_template=context_refine_template,)def_get_nodes(self,message:str)->List[NodeWithScore]:"""Generate context information from a message."""nodes=self._retriever.retrieve(message)forpostprocessorinself._node_postprocessors:nodes=postprocessor.postprocess_nodes(nodes,query_bundle=QueryBundle(message))returnnodesasyncdef_aget_nodes(self,message:str)->List[NodeWithScore]:"""Generate context information from a message."""nodes=awaitself._retriever.aretrieve(message)forpostprocessorinself._node_postprocessors:nodes=postprocessor.postprocess_nodes(nodes,query_bundle=QueryBundle(message))returnnodesdef_get_response_synthesizer(self,chat_history:List[ChatMessage],streaming:bool=False)->CompactAndRefine:# Pull the system prompt from the prefix messagessystem_prompt=""prefix_messages=self._prefix_messagesif(len(self._prefix_messages)!=0andself._prefix_messages[0].role==MessageRole.SYSTEM):system_prompt=str(self._prefix_messages[0].content)prefix_messages=self._prefix_messages[1:]# Get the messages for the QA and refine promptsqa_messages=get_prefix_messages_with_context(self._context_template,system_prompt,prefix_messages,chat_history,self._llm.metadata.system_role,)refine_messages=get_prefix_messages_with_context(self._context_refine_template,system_prompt,prefix_messages,chat_history,self._llm.metadata.system_role,)# Get the response synthesizerreturnget_response_synthesizer(self._llm,self.callback_manager,qa_messages,refine_messages,streaming,qa_function_mappings=self._context_template.function_mappings,refine_function_mappings=self._context_refine_template.function_mappings,)@trace_method("chat")defchat(self,message:str,chat_history:Optional[List[ChatMessage]]=None,prev_chunks:Optional[List[NodeWithScore]]=None,)->AgentChatResponse:ifchat_historyisnotNone:self._memory.set(chat_history)# get nodes and postprocess themnodes=self._get_nodes(message)iflen(nodes)==0andprev_chunksisnotNone:nodes=prev_chunks# Get the response synthesizer with dynamic promptschat_history=self._memory.get(input=message,)synthesizer=self._get_response_synthesizer(chat_history)response=synthesizer.synthesize(message,nodes)user_message=ChatMessage(content=message,role=MessageRole.USER)ai_message=ChatMessage(content=str(response),role=MessageRole.ASSISTANT)self._memory.put(user_message)self._memory.put(ai_message)returnAgentChatResponse(response=str(response),sources=[ToolOutput(tool_name="retriever",content=str(nodes),raw_input={"message":message},raw_output=nodes,)],source_nodes=nodes,)@trace_method("chat")defstream_chat(self,message:str,chat_history:Optional[List[ChatMessage]]=None,prev_chunks:Optional[List[NodeWithScore]]=None,)->StreamingAgentChatResponse:ifchat_historyisnotNone:self._memory.set(chat_history)# get nodes and postprocess themnodes=self._get_nodes(message)iflen(nodes)==0andprev_chunksisnotNone:nodes=prev_chunks# Get the response synthesizer with dynamic promptschat_history=self._memory.get(input=message,)synthesizer=self._get_response_synthesizer(chat_history,streaming=True)response=synthesizer.synthesize(message,nodes)assertisinstance(response,StreamingResponse)defwrapped_gen(response:StreamingResponse)->ChatResponseGen:full_response=""fortokeninresponse.response_gen:full_response+=tokenyieldChatResponse(message=ChatMessage(content=full_response,role=MessageRole.ASSISTANT),delta=token,)user_message=ChatMessage(content=message,role=MessageRole.USER)ai_message=ChatMessage(content=full_response,role=MessageRole.ASSISTANT)self._memory.put(user_message)self._memory.put(ai_message)returnStreamingAgentChatResponse(chat_stream=wrapped_gen(response),sources=[ToolOutput(tool_name="retriever",content=str(nodes),raw_input={"message":message},raw_output=nodes,)],source_nodes=nodes,is_writing_to_memory=False,)@trace_method("chat")asyncdefachat(self,message:str,chat_history:Optional[List[ChatMessage]]=None,prev_chunks:Optional[List[NodeWithScore]]=None,)->AgentChatResponse:ifchat_historyisnotNone:self._memory.set(chat_history)# get nodes and postprocess themnodes=awaitself._aget_nodes(message)iflen(nodes)==0andprev_chunksisnotNone:nodes=prev_chunks# Get the response synthesizer with dynamic promptschat_history=self._memory.get(input=message,)synthesizer=self._get_response_synthesizer(chat_history)response=awaitsynthesizer.asynthesize(message,nodes)user_message=ChatMessage(content=message,role=MessageRole.USER)ai_message=ChatMessage(content=str(response),role=MessageRole.ASSISTANT)awaitself._memory.aput(user_message)awaitself._memory.aput(ai_message)returnAgentChatResponse(response=str(response),sources=[ToolOutput(tool_name="retriever",content=str(nodes),raw_input={"message":message},raw_output=nodes,)],source_nodes=nodes,)@trace_method("chat")asyncdefastream_chat(self,message:str,chat_history:Optional[List[ChatMessage]]=None,prev_chunks:Optional[List[NodeWithScore]]=None,)->StreamingAgentChatResponse:ifchat_historyisnotNone:self._memory.set(chat_history)# get nodes and postprocess themnodes=awaitself._aget_nodes(message)iflen(nodes)==0andprev_chunksisnotNone:nodes=prev_chunks# Get the response synthesizer with dynamic promptschat_history=self._memory.get(input=message,)synthesizer=self._get_response_synthesizer(chat_history,streaming=True)response=awaitsynthesizer.asynthesize(message,nodes)assertisinstance(response,AsyncStreamingResponse)asyncdefwrapped_gen(response:AsyncStreamingResponse)->ChatResponseAsyncGen:full_response=""asyncfortokeninresponse.async_response_gen():full_response+=tokenyieldChatResponse(message=ChatMessage(content=full_response,role=MessageRole.ASSISTANT),delta=token,)user_message=ChatMessage(content=message,role=MessageRole.USER)ai_message=ChatMessage(content=full_response,role=MessageRole.ASSISTANT)awaitself._memory.aput(user_message)awaitself._memory.aput(ai_message)returnStreamingAgentChatResponse(achat_stream=wrapped_gen(response),sources=[ToolOutput(tool_name="retriever",content=str(nodes),raw_input={"message":message},raw_output=nodes,)],source_nodes=nodes,is_writing_to_memory=False,)defreset(self)->None:self._memory.reset()@propertydefchat_history(self)->List[ChatMessage]:"""Get chat history."""returnself._memory.get_all()
@classmethoddeffrom_defaults(cls,retriever:BaseRetriever,chat_history:Optional[List[ChatMessage]]=None,memory:Optional[BaseMemory]=None,system_prompt:Optional[str]=None,prefix_messages:Optional[List[ChatMessage]]=None,node_postprocessors:Optional[List[BaseNodePostprocessor]]=None,context_template:Optional[Union[str,PromptTemplate]]=None,context_refine_template:Optional[Union[str,PromptTemplate]]=None,llm:Optional[LLM]=None,**kwargs:Any,)->"ContextChatEngine":"""Initialize a ContextChatEngine from default parameters."""llm=llmorSettings.llmchat_history=chat_historyor[]memory=memoryorChatMemoryBuffer.from_defaults(chat_history=chat_history,token_limit=llm.metadata.context_window-256)ifsystem_promptisnotNone:ifprefix_messagesisnotNone:raiseValueError("Cannot specify both system_prompt and prefix_messages")prefix_messages=[ChatMessage(content=system_prompt,role=llm.metadata.system_role)]prefix_messages=prefix_messagesor[]node_postprocessors=node_postprocessorsor[]returncls(retriever,llm=llm,memory=memory,prefix_messages=prefix_messages,node_postprocessors=node_postprocessors,callback_manager=Settings.callback_manager,context_template=context_template,context_refine_template=context_refine_template,)