From 9799a1b64f5786ab21a82a67822398fc4618e6aa Mon Sep 17 00:00:00 2001 From: linked-liszt Date: Thu, 16 Jan 2025 15:43:02 -0600 Subject: [PATCH 1/2] Add autogen argo LLM --- autogen_llm.py | 74 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 autogen_llm.py diff --git a/autogen_llm.py b/autogen_llm.py new file mode 100644 index 0000000..5eeba6a --- /dev/null +++ b/autogen_llm.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace +import requests +import autogen +from autogen import AssistantAgent, UserProxyAgent +import params + +with open(params.anl_llm_url_path, 'r') as url_f: + ANL_URL = url_f.read().strip() + +class ArgoModelClient: + def __init__(self, config, **kwargs): + print(f"Argo config: {config}") + self.model = config['model'] + self.temp = config['temp'] + + def create(self, params): + if params.get("stream", False) and "messages" in params: + raise NotImplementedError("Local models do not support streaming.") + + response = SimpleNamespace() + + prompt = apply_chat_template(params['messages']) + + response.choices = [] + response.model = "model_name" + + num_of_responses = params.get("n", 1) + + + for _ in range(num_of_responses): + text = self._query_argo(prompt) + choice = SimpleNamespace() + choice.message = SimpleNamespace() + choice.message.content = text + choice.message.function_call = None + response.choices.append(choice) + return response + + def _query_argo(self, prompt): + req_obj = {'user': 'aps', + 'model': self.model, + 'prompt': [prompt], + 'system': "", + 'temperature': self.temp} + result = requests.post(ANL_URL, json=req_obj) + if not result.ok: + raise ValueError(f"error {result.status_code} ({result.reason})") + + response = result.json()['response'] + return response + + def message_retrieval(self, response): + choices = response.choices + return [choice.message.content for choice in choices] + + def cost(self, response) -> float: + response.cost = 0 + return 0 + + @staticmethod + def get_usage(response): + return {} + + +def apply_chat_template(messages): + output_str = "" + for message in messages: + output_str += f'-- {message['role']} ' + if 'name' in message: + output_str += f'{message['name']} --\n' + else: + output_str += '--\n' + output_str += f'{message['content']}\n --- \n\n' + return output_str \ No newline at end of file From 8a36a2422d7bac2c0c171bed928ab75fa3b66897 Mon Sep 17 00:00:00 2001 From: linked-liszt Date: Thu, 16 Jan 2025 16:21:06 -0600 Subject: [PATCH 2/2] Add cli and fix autogen_llm --- autogen_llm.py | 1 + chatbot_cli.py | 245 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 246 insertions(+) create mode 100644 chatbot_cli.py diff --git a/autogen_llm.py b/autogen_llm.py index 5eeba6a..e1e2465 100644 --- a/autogen_llm.py +++ b/autogen_llm.py @@ -41,6 +41,7 @@ def _query_argo(self, prompt): 'model': self.model, 'prompt': [prompt], 'system': "", + 'stop': [], 'temperature': self.temp} result = requests.post(ANL_URL, json=req_obj) if not result.ok: diff --git a/chatbot_cli.py b/chatbot_cli.py new file mode 100644 index 0000000..bcbde0e --- /dev/null +++ b/chatbot_cli.py @@ -0,0 +1,245 @@ +import params +import llms +import bot_tools +from langchain_core.prompts import PromptTemplate +from langchain.chains import LLMChain +from langchain.chains.conversation.memory import ConversationBufferWindowMemory +from langchain.agents import AgentExecutor, create_json_chat_agent +from langchain_core.messages import AIMessage, HumanMessage +import prompts + +class Chat(): + def __init__(self, llm, embedding, doc_store): + self.llm = llm + self.embedding = embedding + self.doc_store = doc_store + self.is_PDF = False #Flag to use NER over right set of docs. Changed in update_pdf_docstore + + + def _init_chain(self): + template = """The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. + +Context: +{context} + +Current conversation: +{history} +Human: {input} +AI:""" + + PROMPT = PromptTemplate( + input_variables=["history", "input", "context"], template=template + ) + memory = ConversationBufferWindowMemory(memory_key="history", + input_key = "input", + k=6) + + conversation = LLMChain( + prompt=PROMPT, + llm=self.llm, + verbose=True, + memory=memory + ) + + return conversation + + + #Method to find text with highest likely context + def _get_context(self, query, doc_store): + + # Context retrieval from embeddings + docs = doc_store.similarity_search_with_score(query, k=params.N_hits) + #Get context strings + context="" + print ("Context hits found", len(docs)) + for i in range(min(params.N_hits, len(docs))): + if docs[i][1] some length from query + ner_hits = [] + + #Set path from where to get NER context hits + if self.is_PDF: + doc_path = params.pdf_text_path + print("Getting NER hits from PDF context") + else: + doc_path = params.doc_path_root + clean_pdf_paths() #Make sure PDF folders are clean to avoid context leak + print("Getting NER hits from facility context") + + for ner in ners: #Grep NEs from raw text + try: + hit = subprocess.check_output("grep -r -i -h '%s' %s/" %(ner, doc_path), shell=True).decode() + hits = hit.split("\n") #split all the grep results into indiv strings + ner_hits.extend(hits) + except subprocess.CalledProcessError as err: + if err.returncode > 1: + print ("No hits found for: ", ner) + continue + #Exit values: 0 One or more lines were selected. 1 No lines were selected. >1 An error occurred. + #print ("NERs", ner_hits) + + ner_hits.sort(key=len, reverse=True) #Sort by length of hits + #print ("Sorted NERs", ner_hits) + + for i in range(min(params.N_NER_hits, len(ner_hits))): + print ("Selected NER hit %d : " %i, ner_hits[i]) + context += ner_hits[i] + + return context + + + def generate_response(self, history, debug_output, convo_state, doc_state = None): + user_message = history[-1]['content'] #History is list of tuple list. E.g. : [['Hi', 'Test'], ['Hello again', '']] + all_user_messages = [x['content'] for x in history] + + if convo_state is None: + convo_state = self._init_chain() + + if self.doc_store is not None: + context = "" + for message in all_user_messages: + context += self._get_context(message, self.doc_store) + elif doc_state is not None: + context = "" + for message in all_user_messages: + context += self._get_context(message, doc_state) + else: + context = "" + + if debug_output: + inputs = convo_state.prep_inputs({'input': user_message, 'context':context}) + prompt = convo_state.prep_prompts([inputs])[0][0].text + + bot_message = convo_state.predict(input=user_message, context=context) + + + if debug_output: + bot_message = f'---Prompt---\n\n {prompt} \n\n---Response---\n\n {bot_message}' + + print(history) + print(convo_state) + history.append( + ChatMessage(role='assistant', content=bot_message) + ) + + return history, convo_state + + def add_message(self, user_message, history): + history.append( + ChatMessage(role='user', content=user_message) + ) + return "", history + + def clear_memory(self, convo_state): + if convo_state is not None: + convo_state.memory.clear() + return convo_state, None + else: + return None, None + +class ToolChat(Chat): + """ + Implements an agentexector in a chat context. The agentexecutor is called in a fundimentally + differnet way than the other chains, so custom implementaiton for much of the class. + """ + def _init_chain(self): + """ + tools = [ + dfrac_tools.DiffractometerAIO(params.spec_init) + ] + """ + # TODO: CHANGE CREATION TYPE + tools = [bot_tools.lattice_tool, bot_tools.diffractometer_tool] + + memory = ConversationBufferWindowMemory(memory_key="chat_history", k=6) + agent = create_json_chat_agent( + tools=tools, + llm=self.llm, + prompt=prompts.json_tool_prompt) + + agent_executor = AgentExecutor( + agent=agent, tools=tools, handle_parsing_errors=True, + max_iterations = 15, + verbose=True + ) + + self.memory = memory + self.conversation = agent_executor + + return memory, agent_executor + + def generate_response(self, history, debug_output): + user_message = history[-1]['content'] #History is list of tuple list. E.g. : [['Hi', 'Test'], ['Hello again', '']] + + # Convert to langchain history + lang_hist = [] + for message in history: + if message['role'] == 'user': + lang_hist.append(HumanMessage(content=message['content'])) + elif message['role'] == 'assistant': + lang_hist.append(AIMessage(content=message['content'])) + else: + raise ValueError(f"Unknown role in history {history}, {message['role']}. Add way to resolve.") + + #raise ValueError(f'Unknown role in history {history}, {message['role']}. Add way to resolve.') + + # TODO: Implement debug output for langchain agents. Might have to use a callback? + print(f'User input: {user_message}') + response = self.conversation.invoke( + { + "input": user_message, + "chat_history": lang_hist, + } + ) + + bot_message = response['output'] + #Pass user message and get context and pass to model + history.append( + {'role':'assistant', 'content':bot_message} + ) + + return history + +class PolybotExecChat(ToolChat): + def _init_chain(self): + tools = [bot_tools.exec_polybot_tool, bot_tools.exec_polybot_lint_tool] + + memory = ConversationBufferWindowMemory(memory_key="chat_history", k=7) + + + agent = create_json_chat_agent( + tools=tools, + llm=self.llm, + prompt=prompts.json_tool_prompt) + + agent_executor = AgentExecutor( + agent=agent, tools=tools, handle_parsing_errors=True, + max_iterations = 15, + verbose=True + ) + + self.memory = memory + self.conversation = agent_executor + + return memory, agent_executor + +llm = llms.AnlLLM(params) +embeddings = llms.ANLEmbeddingModel(params) + + +polybot_exec = PolybotExecChat(llm, embeddings, None) +polybot_exec._init_chain() + + +chat_history = [{'content': 'Pick up the vial in rack 1', 'role': 'user'}] +#Pass an empty string to context when don't want domain specific context +chat_history = polybot_exec.generate_response(chat_history, True) + +print(chat_history[-1]) +