team-ai / chains.py
peichao.dong
add rewrite context
a72034e
raw
history blame
1.71 kB
from typing import Any, Optional
from langchain.chains import LLMChain
from langchain.base_language import BaseLanguageModel
from langchain.schema import LLMResult, PromptValue
from langchain.prompts import PromptTemplate
from langchain.memory.chat_memory import BaseMemory
from langchain.chat_models import ChatOpenAI
from promopts import CONTENT_RE_WRIGHT_PROMPT, FEEDBACK_PROMPT
class HumanFeedBackChain(LLMChain):
"""Chain to run queries against LLMs."""
memory: Optional[BaseMemory] = None
def __init__(self, verbose=True, llm: BaseLanguageModel = ChatOpenAI(temperature=0.7), memory: Optional[BaseMemory] = None, prompt: PromptTemplate = FEEDBACK_PROMPT):
super().__init__(llm=llm, prompt=prompt, memory=memory, verbose=verbose)
def run(self, *args: Any, **kwargs: Any) -> str:
"""Run the chain as text in, text out or multiple variables, text out."""
if len(self.output_keys) != 1:
raise ValueError(
f"`run` not supported when there is not exactly "
f"one output key. Got {self.output_keys}."
)
if args and not kwargs:
if len(args) != 1:
raise ValueError(
"`run` supports only one positional argument.")
return self("Answer:" + args[0])[self.output_keys[0]]
if kwargs and not args:
return self(kwargs)[self.output_keys[0]]
raise ValueError(
f"`run` supported with either positional arguments or keyword arguments"
f" but not both. Got args: {args} and kwargs: {kwargs}."
)
contextRewriteChain = LLMChain(llm=ChatOpenAI(temperature=0.7), prompt=CONTENT_RE_WRIGHT_PROMPT)