diviner/service/agent/rag.py

# -*- coding: utf-8 -*-
from .__load__ import *
from langchain_core.language_models.chat_models import HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

class Chat(object):

    chain = False
    def say(self, content):
        return [HumanMessage(content=content)]

    def set(self, prompts = []):
        chain = ChatPromptTemplate.from_messages(prompts)
        if not self.chain:
            self.chain = chain
        else:
            self.chain = self.chain | chain

    def out(self, handle, var = {}):
        chain = self.chain | handle | StrOutputParser()
        for chunk in chain.stream(var):
            print(chunk, end="")