File size: 6,361 Bytes
7c24f26 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
"""
Langchain agent
"""
from typing import Generator, Dict, Optional, Literal, TypedDict, List
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain.memory import ConversationBufferMemory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import BaseMessage
from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableSerializable
from langchain_core.output_parsers import StrOutputParser
from .prompts import SYSTEM_PROMPT, REFERENCE_SYSTEM_PROMPT
load_dotenv()
valid_model_names = Literal[
'llama3-70b-8192',
'llama3-8b-8192',
'gemma-7b-it',
'gemma2-9b-it',
'mixtral-8x7b-32768'
]
class ResponseChunk(TypedDict):
delta: str
response_type: Literal['intermediate', 'output']
metadata: Dict = {}
class MOAgent:
def __init__(
self,
main_agent: RunnableSerializable[Dict, str],
layer_agent: RunnableSerializable[Dict, Dict],
reference_system_prompt: Optional[str] = None,
cycles: Optional[int] = None,
chat_memory: Optional[ConversationBufferMemory] = None
) -> None:
self.reference_system_prompt = reference_system_prompt or REFERENCE_SYSTEM_PROMPT
self.main_agent = main_agent
self.layer_agent = layer_agent
self.cycles = cycles or 1
self.chat_memory = chat_memory or ConversationBufferMemory(
memory_key="messages",
return_messages=True
)
@staticmethod
def concat_response(
inputs: Dict[str, str],
reference_system_prompt: Optional[str] = None
):
reference_system_prompt = reference_system_prompt or REFERENCE_SYSTEM_PROMPT
responses = ""
res_list = []
for i, out in enumerate(inputs.values()):
responses += f"{i}. {out}\n"
res_list.append(out)
formatted_prompt = reference_system_prompt.format(responses=responses)
return {
'formatted_response': formatted_prompt,
'responses': res_list
}
@classmethod
def from_config(
cls,
main_model: Optional[valid_model_names] = 'llama3-70b-8192',
system_prompt: Optional[str] = None,
cycles: int = 1,
layer_agent_config: Optional[Dict] = None,
reference_system_prompt: Optional[str] = None,
**main_model_kwargs
):
reference_system_prompt = reference_system_prompt or REFERENCE_SYSTEM_PROMPT
system_prompt = system_prompt or SYSTEM_PROMPT
layer_agent = MOAgent._configure_layer_agent(layer_agent_config)
main_agent = MOAgent._create_agent_from_system_prompt(
system_prompt=system_prompt,
model_name=main_model,
**main_model_kwargs
)
return cls(
main_agent=main_agent,
layer_agent=layer_agent,
reference_system_prompt=reference_system_prompt,
cycles=cycles
)
@staticmethod
def _configure_layer_agent(
layer_agent_config: Optional[Dict] = None
) -> RunnableSerializable[Dict, Dict]:
if not layer_agent_config:
layer_agent_config = {
'layer_agent_1' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'llama3-8b-8192'},
'layer_agent_2' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'gemma-7b-it'},
'layer_agent_3' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'mixtral-8x7b-32768'}
}
parallel_chain_map = dict()
for key, value in layer_agent_config.items():
chain = MOAgent._create_agent_from_system_prompt(
system_prompt=value.pop("system_prompt", SYSTEM_PROMPT),
model_name=value.pop("model_name", 'llama3-8b-8192'),
**value
)
parallel_chain_map[key] = RunnablePassthrough() | chain
chain = parallel_chain_map | RunnableLambda(MOAgent.concat_response)
return chain
@staticmethod
def _create_agent_from_system_prompt(
system_prompt: str = SYSTEM_PROMPT,
model_name: str = "llama3-8b-8192",
**llm_kwargs
) -> RunnableSerializable[Dict, str]:
prompt = ChatPromptTemplate.from_messages([
("system", system_prompt),
MessagesPlaceholder(variable_name="messages", optional=True),
("human", "{input}")
])
assert 'helper_response' in prompt.input_variables
llm = ChatGroq(model=model_name, **llm_kwargs)
chain = prompt | llm | StrOutputParser()
return chain
def chat(
self,
input: str,
messages: Optional[List[BaseMessage]] = None,
cycles: Optional[int] = None,
save: bool = True,
output_format: Literal['string', 'json'] = 'string'
) -> Generator[str | ResponseChunk, None, None]:
cycles = cycles or self.cycles
llm_inp = {
'input': input,
'messages': messages or self.chat_memory.load_memory_variables({})['messages'],
'helper_response': ""
}
for cyc in range(cycles):
layer_output = self.layer_agent.invoke(llm_inp)
l_frm_resp = layer_output['formatted_response']
l_resps = layer_output['responses']
llm_inp = {
'input': input,
'messages': self.chat_memory.load_memory_variables({})['messages'],
'helper_response': l_frm_resp
}
if output_format == 'json':
for l_out in l_resps:
yield ResponseChunk(
delta=l_out,
response_type='intermediate',
metadata={'layer': cyc + 1}
)
stream = self.main_agent.stream(llm_inp)
response = ""
for chunk in stream:
if output_format == 'json':
yield ResponseChunk(
delta=chunk,
response_type='output',
metadata={}
)
else:
yield chunk
response += chunk
if save:
self.chat_memory.save_context({'input': input}, {'output': response}) |