From 8935dc4e313b98db7eb19ea9d08255159eaff2aa Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 24 Dec 2024 00:52:14 +0000 Subject: [PATCH 01/14] feat: add tool parsing --- application/llm/google_ai.py | 93 +++++++++++++++++++++++++++++++++--- 1 file changed, 87 insertions(+), 6 deletions(-) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index df252abfc..33ae0855f 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -1,11 +1,13 @@ from application.llm.base import BaseLLM +from application.core.settings import settings +import logging class GoogleLLM(BaseLLM): def __init__(self, api_key=None, user_api_key=None, *args, **kwargs): super().__init__(*args, **kwargs) - self.api_key = api_key + self.api_key = settings.API_KEY self.user_api_key = user_api_key def _clean_messages_google(self, messages): @@ -16,6 +18,52 @@ def _clean_messages_google(self, messages): } for message in messages[1:] ] + + def _clean_tools_format(self, tools_data): + """ + Cleans the tools data format, converting string type representations + to the expected dictionary structure for google-generativeai. + """ + if isinstance(tools_data, list): + return [self._clean_tools_format(item) for item in tools_data] + elif isinstance(tools_data, dict): + if 'function' in tools_data and 'type' in tools_data and tools_data['type'] == 'function': + # Handle the case where tools are nested under 'function' + cleaned_function = self._clean_tools_format(tools_data['function']) + return {'function_declarations': [cleaned_function]} + elif 'function' in tools_data and 'type_' in tools_data and tools_data['type_'] == 'function': + # Handle the case where tools are nested under 'function' and type is already 'type_' + cleaned_function = self._clean_tools_format(tools_data['function']) + return {'function_declarations': [cleaned_function]} + else: + new_tools_data = {} + for key, value in tools_data.items(): + if key == 'type': + if value == 'string': + new_tools_data['type_'] = 'STRING' # Keep as string for now + elif value == 'object': + new_tools_data['type_'] = 'OBJECT' # Keep as string for now + elif key == 'additionalProperties': + continue + elif key == 'properties': + if isinstance(value, dict): + new_properties = {} + for prop_name, prop_value in value.items(): + if isinstance(prop_value, dict) and 'type' in prop_value: + if prop_value['type'] == 'string': + new_properties[prop_name] = {'type_': 'STRING', 'description': prop_value.get('description')} + # Add more type mappings as needed + else: + new_properties[prop_name] = self._clean_tools_format(prop_value) + new_tools_data[key] = new_properties + else: + new_tools_data[key] = self._clean_tools_format(value) + + else: + new_tools_data[key] = self._clean_tools_format(value) + return new_tools_data + else: + return tools_data def _raw_gen( self, @@ -23,12 +71,29 @@ def _raw_gen( model, messages, stream=False, + tools=None, **kwargs ): import google.generativeai as genai genai.configure(api_key=self.api_key) - model = genai.GenerativeModel(model, system_instruction=messages[0]["content"]) - response = model.generate_content(self._clean_messages_google(messages)) + + config = { + } + model = 'gemini-2.0-flash-exp' + + model = genai.GenerativeModel( + model_name=model, + generation_config=config, + system_instruction=messages[0]["content"], + tools=self._clean_tools_format(tools) + ) + chat_session = model.start_chat( + history=self._clean_messages_google(messages)[:-1] + ) + response = chat_session.send_message( + self._clean_messages_google(messages)[-1] + ) + logging.info(response) return response.text def _raw_gen_stream( @@ -37,12 +102,28 @@ def _raw_gen_stream( model, messages, stream=True, + tools=None, **kwargs ): import google.generativeai as genai genai.configure(api_key=self.api_key) - model = genai.GenerativeModel(model, system_instruction=messages[0]["content"]) - response = model.generate_content(self._clean_messages_google(messages), stream=True) + config = { + } + model = genai.GenerativeModel( + model_name=model, + generation_config=config, + system_instruction=messages[0]["content"] + ) + chat_session = model.start_chat( + history=self._clean_messages_google(messages)[:-1], + ) + response = chat_session.send_message( + self._clean_messages_google(messages)[-1] + , stream=stream + ) for line in response: if line.text is not None: - yield line.text \ No newline at end of file + yield line.text + + def _supports_tools(self): + return True \ No newline at end of file From 51225b18b2e520472d0f969464925dea8b5c4fc6 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 13 Jan 2025 10:37:53 +0000 Subject: [PATCH 02/14] add google --- application/llm/google_ai.py | 39 ++++++++++++++++++++++-------------- application/tools/agent.py | 32 ++++++++++++++++++++++++++++- 2 files changed, 55 insertions(+), 16 deletions(-) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 33ae0855f..09f0d9f70 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -72,29 +72,38 @@ def _raw_gen( messages, stream=False, tools=None, + formatting="openai", **kwargs ): - import google.generativeai as genai - genai.configure(api_key=self.api_key) + from google import genai + from google.genai import types + client = genai.Client(api_key=self.api_key) + config = { } model = 'gemini-2.0-flash-exp' + if formatting=="raw": + response = client.models.generate_content( + model=model, + contents=messages + ) - model = genai.GenerativeModel( - model_name=model, - generation_config=config, - system_instruction=messages[0]["content"], - tools=self._clean_tools_format(tools) + else: + model = genai.GenerativeModel( + model_name=model, + generation_config=config, + system_instruction=messages[0]["content"], + tools=self._clean_tools_format(tools) + ) + chat_session = model.start_chat( + history=self._clean_messages_google(messages)[:-1] ) - chat_session = model.start_chat( - history=self._clean_messages_google(messages)[:-1] - ) - response = chat_session.send_message( - self._clean_messages_google(messages)[-1] - ) - logging.info(response) - return response.text + response = chat_session.send_message( + self._clean_messages_google(messages)[-1] + ) + logging.info(response) + return response.text def _raw_gen_stream( self, diff --git a/application/tools/agent.py b/application/tools/agent.py index d4077e45d..a6aec2e98 100644 --- a/application/tools/agent.py +++ b/application/tools/agent.py @@ -1,4 +1,5 @@ import json +import logging from application.core.mongo_db import MongoDB from application.llm.llm_creator import LLMCreator @@ -79,6 +80,25 @@ def _execute_tool_action(self, tools_dict, call): print(f"Executing tool: {action_name} with args: {call_args}") return tool.execute_action(action_name, **call_args), call_id + def _execute_tool_action_google(self, tools_dict, call): + call_args = json.loads(call.args) + tool_id = call.name.split("_")[-1] + action_name = call.name.rsplit("_", 1)[0] + + tool_data = tools_dict[tool_id] + action_data = next( + action for action in tool_data["actions"] if action["name"] == action_name + ) + + for param, details in action_data["parameters"]["properties"].items(): + if param not in call_args and "value" in details: + call_args[param] = details["value"] + + tm = ToolManager(config={}) + tool = tm.load_tool(tool_data["name"], tool_config=tool_data["config"]) + print(f"Executing tool: {action_name} with args: {call_args}") + return tool.execute_action(action_name, **call_args) + def _simple_tool_agent(self, messages): tools_dict = self._get_user_tools() self._prepare_tools(tools_dict) @@ -91,8 +111,18 @@ def _simple_tool_agent(self, messages): if resp.message.content: yield resp.message.content return + # check if self.llm class is GoogleLLM + while self.llm.__class__.__name__ == "GoogleLLM" and resp.content.parts[0].function_call: + from google.genai import types + + function_call_part = resp.candidates[0].content.parts[0] + tool_response = self._execute_tool_action_google(tools_dict, function_call_part.function_call) + function_response_part = types.Part.from_function_response( + name=function_call_part.function_call.name, + response=tool_response + ) - while resp.finish_reason == "tool_calls": + while self.llm.__class__.__name__ == "OpenAILLM" and resp.finish_reason == "tool_calls": message = json.loads(resp.model_dump_json())["message"] keys_to_remove = {"audio", "function_call", "refusal"} filtered_data = { From 811dfecf9888ce8722563f9315897b9862371099 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Wed, 15 Jan 2025 16:35:26 +0530 Subject: [PATCH 03/14] refactor: tool agent for action parser and handlers --- application/llm/openai.py | 20 +++---- application/tools/agent.py | 76 +++---------------------- application/tools/llm_handler.py | 74 ++++++++++++++++++++++++ application/tools/tool_action_parser.py | 26 +++++++++ 4 files changed, 116 insertions(+), 80 deletions(-) create mode 100644 application/tools/llm_handler.py create mode 100644 application/tools/tool_action_parser.py diff --git a/application/llm/openai.py b/application/llm/openai.py index cc2285a12..b507a1da8 100644 --- a/application/llm/openai.py +++ b/application/llm/openai.py @@ -1,6 +1,5 @@ -from application.llm.base import BaseLLM from application.core.settings import settings - +from application.llm.base import BaseLLM class OpenAILLM(BaseLLM): @@ -10,10 +9,7 @@ def __init__(self, api_key=None, user_api_key=None, *args, **kwargs): super().__init__(*args, **kwargs) if settings.OPENAI_BASE_URL: - self.client = OpenAI( - api_key=api_key, - base_url=settings.OPENAI_BASE_URL - ) + self.client = OpenAI(api_key=api_key, base_url=settings.OPENAI_BASE_URL) else: self.client = OpenAI(api_key=api_key) self.api_key = api_key @@ -27,8 +23,8 @@ def _raw_gen( stream=False, tools=None, engine=settings.AZURE_DEPLOYMENT_NAME, - **kwargs - ): + **kwargs, + ): if tools: response = self.client.chat.completions.create( model=model, messages=messages, stream=stream, tools=tools, **kwargs @@ -48,18 +44,16 @@ def _raw_gen_stream( stream=True, tools=None, engine=settings.AZURE_DEPLOYMENT_NAME, - **kwargs - ): + **kwargs, + ): response = self.client.chat.completions.create( model=model, messages=messages, stream=stream, **kwargs ) for line in response: - # import sys - # print(line.choices[0].delta.content, file=sys.stderr) if line.choices[0].delta.content is not None: yield line.choices[0].delta.content - + def _supports_tools(self): return True diff --git a/application/tools/agent.py b/application/tools/agent.py index a6aec2e98..f4b37d9bf 100644 --- a/application/tools/agent.py +++ b/application/tools/agent.py @@ -1,8 +1,7 @@ -import json -import logging - from application.core.mongo_db import MongoDB from application.llm.llm_creator import LLMCreator +from application.tools.llm_handler import get_llm_handler +from application.tools.tool_action_parser import ToolActionParser from application.tools.tool_manager import ToolManager @@ -12,6 +11,7 @@ def __init__(self, llm_name, gpt_model, api_key, user_api_key=None): self.llm = LLMCreator.create_llm( llm_name, api_key=api_key, user_api_key=user_api_key ) + self.llm_handler = get_llm_handler(llm_name) self.gpt_model = gpt_model # Static tool configuration (to be replaced later) self.tools = [] @@ -61,29 +61,8 @@ def _prepare_tools(self, tools_dict): ] def _execute_tool_action(self, tools_dict, call): - call_id = call.id - call_args = json.loads(call.function.arguments) - tool_id = call.function.name.split("_")[-1] - action_name = call.function.name.rsplit("_", 1)[0] - - tool_data = tools_dict[tool_id] - action_data = next( - action for action in tool_data["actions"] if action["name"] == action_name - ) - - for param, details in action_data["parameters"]["properties"].items(): - if param not in call_args and "value" in details: - call_args[param] = details["value"] - - tm = ToolManager(config={}) - tool = tm.load_tool(tool_data["name"], tool_config=tool_data["config"]) - print(f"Executing tool: {action_name} with args: {call_args}") - return tool.execute_action(action_name, **call_args), call_id - - def _execute_tool_action_google(self, tools_dict, call): - call_args = json.loads(call.args) - tool_id = call.name.split("_")[-1] - action_name = call.name.rsplit("_", 1)[0] + parser = ToolActionParser(self.llm.__class__.__name__) + tool_id, action_name, call_args = parser.parse_args(call) tool_data = tools_dict[tool_id] action_data = next( @@ -97,7 +76,9 @@ def _execute_tool_action_google(self, tools_dict, call): tm = ToolManager(config={}) tool = tm.load_tool(tool_data["name"], tool_config=tool_data["config"]) print(f"Executing tool: {action_name} with args: {call_args}") - return tool.execute_action(action_name, **call_args) + result = tool.execute_action(action_name, **call_args) + call_id = getattr(call, "id", None) + return result, call_id def _simple_tool_agent(self, messages): tools_dict = self._get_user_tools() @@ -111,47 +92,8 @@ def _simple_tool_agent(self, messages): if resp.message.content: yield resp.message.content return - # check if self.llm class is GoogleLLM - while self.llm.__class__.__name__ == "GoogleLLM" and resp.content.parts[0].function_call: - from google.genai import types - function_call_part = resp.candidates[0].content.parts[0] - tool_response = self._execute_tool_action_google(tools_dict, function_call_part.function_call) - function_response_part = types.Part.from_function_response( - name=function_call_part.function_call.name, - response=tool_response - ) - - while self.llm.__class__.__name__ == "OpenAILLM" and resp.finish_reason == "tool_calls": - message = json.loads(resp.model_dump_json())["message"] - keys_to_remove = {"audio", "function_call", "refusal"} - filtered_data = { - k: v for k, v in message.items() if k not in keys_to_remove - } - messages.append(filtered_data) - tool_calls = resp.message.tool_calls - for call in tool_calls: - try: - tool_response, call_id = self._execute_tool_action(tools_dict, call) - messages.append( - { - "role": "tool", - "content": str(tool_response), - "tool_call_id": call_id, - } - ) - except Exception as e: - messages.append( - { - "role": "tool", - "content": f"Error executing tool: {str(e)}", - "tool_call_id": call.id, - } - ) - # Generate a new response from the LLM after processing tools - resp = self.llm.gen( - model=self.gpt_model, messages=messages, tools=self.tools - ) + resp = self.llm_handler.handle_response(self, resp, tools_dict, messages) # If no tool calls are needed, generate the final response if isinstance(resp, str): diff --git a/application/tools/llm_handler.py b/application/tools/llm_handler.py new file mode 100644 index 000000000..58fce56e4 --- /dev/null +++ b/application/tools/llm_handler.py @@ -0,0 +1,74 @@ +import json +from abc import ABC, abstractmethod + + +class LLMHandler(ABC): + @abstractmethod + def handle_response(self, agent, resp, tools_dict, messages, **kwargs): + pass + + +class OpenAILLMHandler(LLMHandler): + def handle_response(self, agent, resp, tools_dict, messages): + while resp.finish_reason == "tool_calls": + message = json.loads(resp.model_dump_json())["message"] + keys_to_remove = {"audio", "function_call", "refusal"} + filtered_data = { + k: v for k, v in message.items() if k not in keys_to_remove + } + messages.append(filtered_data) + + tool_calls = resp.message.tool_calls + for call in tool_calls: + try: + tool_response, call_id = agent._execute_tool_action( + tools_dict, call + ) + messages.append( + { + "role": "tool", + "content": str(tool_response), + "tool_call_id": call_id, + } + ) + except Exception as e: + messages.append( + { + "role": "tool", + "content": f"Error executing tool: {str(e)}", + "tool_call_id": call_id, + } + ) + resp = agent.llm.gen( + model=agent.gpt_model, messages=messages, tools=agent.tools + ) + return resp + + +class GoogleLLMHandler(LLMHandler): + def handle_response(self, agent, resp, tools_dict, messages): + from google.genai import types + + while resp.content.parts[0].function_call: + function_call_part = resp.candidates[0].content.parts[0] + tool_response, call_id = agent._execute_tool_action( + tools_dict, function_call_part.function_call + ) + function_response_part = types.Part.from_function_response( + name=function_call_part.function_call.name, response=tool_response + ) + + messages.append(function_call_part, function_response_part) + resp = agent.llm.gen( + model=agent.gpt_model, messages=messages, tools=agent.tools + ) + + return resp + + +def get_llm_handler(llm_type): + handlers = { + "openai": OpenAILLMHandler(), + "google": GoogleLLMHandler(), + } + return handlers.get(llm_type, OpenAILLMHandler()) diff --git a/application/tools/tool_action_parser.py b/application/tools/tool_action_parser.py new file mode 100644 index 000000000..b708992a5 --- /dev/null +++ b/application/tools/tool_action_parser.py @@ -0,0 +1,26 @@ +import json + + +class ToolActionParser: + def __init__(self, llm_type): + self.llm_type = llm_type + self.parsers = { + "OpenAILLM": self._parse_openai_llm, + "GoogleLLM": self._parse_google_llm, + } + + def parse_args(self, call): + parser = self.parsers.get(self.llm_type, self._parse_openai_llm) + return parser(call) + + def _parse_openai_llm(self, call): + call_args = json.loads(call.function.arguments) + tool_id = call.function.name.split("_")[-1] + action_name = call.function.name.rsplit("_", 1)[0] + return tool_id, action_name, call_args + + def _parse_google_llm(self, call): + call_args = json.loads(call.args) + tool_id = call.name.split("_")[-1] + action_name = call.name.rsplit("_", 1)[0] + return tool_id, action_name, call_args From c97d1e336308c0717bc12774e663d49b403aa345 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Fri, 17 Jan 2025 09:22:41 +0530 Subject: [PATCH 04/14] fix: google parser, llm handler and other errors --- application/llm/google_ai.py | 184 +++++++++++++----------- application/tools/agent.py | 4 +- application/tools/llm_handler.py | 44 ++++-- application/tools/tool_action_parser.py | 9 +- 4 files changed, 140 insertions(+), 101 deletions(-) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 09f0d9f70..24043d9cd 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -1,60 +1,77 @@ -from application.llm.base import BaseLLM +import google.generativeai as genai + from application.core.settings import settings -import logging +from application.llm.base import BaseLLM + class GoogleLLM(BaseLLM): def __init__(self, api_key=None, user_api_key=None, *args, **kwargs): - super().__init__(*args, **kwargs) self.api_key = settings.API_KEY - self.user_api_key = user_api_key + genai.configure(api_key=self.api_key) def _clean_messages_google(self, messages): - return [ - { - "role": "model" if message["role"] == "system" else message["role"], - "parts": [message["content"]], - } - for message in messages[1:] - ] - + cleaned_messages = [] + for message in messages[1:]: + cleaned_messages.append( + { + "role": "model" if message["role"] == "system" else message["role"], + "parts": [message["content"]], + } + ) + return cleaned_messages + def _clean_tools_format(self, tools_data): - """ - Cleans the tools data format, converting string type representations - to the expected dictionary structure for google-generativeai. - """ if isinstance(tools_data, list): return [self._clean_tools_format(item) for item in tools_data] elif isinstance(tools_data, dict): - if 'function' in tools_data and 'type' in tools_data and tools_data['type'] == 'function': + if ( + "function" in tools_data + and "type" in tools_data + and tools_data["type"] == "function" + ): # Handle the case where tools are nested under 'function' - cleaned_function = self._clean_tools_format(tools_data['function']) - return {'function_declarations': [cleaned_function]} - elif 'function' in tools_data and 'type_' in tools_data and tools_data['type_'] == 'function': + cleaned_function = self._clean_tools_format(tools_data["function"]) + return {"function_declarations": [cleaned_function]} + elif ( + "function" in tools_data + and "type_" in tools_data + and tools_data["type_"] == "function" + ): # Handle the case where tools are nested under 'function' and type is already 'type_' - cleaned_function = self._clean_tools_format(tools_data['function']) - return {'function_declarations': [cleaned_function]} + cleaned_function = self._clean_tools_format(tools_data["function"]) + return {"function_declarations": [cleaned_function]} else: new_tools_data = {} for key, value in tools_data.items(): - if key == 'type': - if value == 'string': - new_tools_data['type_'] = 'STRING' # Keep as string for now - elif value == 'object': - new_tools_data['type_'] = 'OBJECT' # Keep as string for now - elif key == 'additionalProperties': + if key == "type": + if value == "string": + new_tools_data["type_"] = "STRING" + elif value == "object": + new_tools_data["type_"] = "OBJECT" + elif key == "additionalProperties": continue - elif key == 'properties': + elif key == "properties": if isinstance(value, dict): new_properties = {} for prop_name, prop_value in value.items(): - if isinstance(prop_value, dict) and 'type' in prop_value: - if prop_value['type'] == 'string': - new_properties[prop_name] = {'type_': 'STRING', 'description': prop_value.get('description')} + if ( + isinstance(prop_value, dict) + and "type" in prop_value + ): + if prop_value["type"] == "string": + new_properties[prop_name] = { + "type_": "STRING", + "description": prop_value.get( + "description" + ), + } # Add more type mappings as needed else: - new_properties[prop_name] = self._clean_tools_format(prop_value) + new_properties[prop_name] = ( + self._clean_tools_format(prop_value) + ) new_tools_data[key] = new_properties else: new_tools_data[key] = self._clean_tools_format(value) @@ -74,65 +91,64 @@ def _raw_gen( tools=None, formatting="openai", **kwargs - ): - from google import genai - from google.genai import types - client = genai.Client(api_key=self.api_key) + ): + config = {} + model_name = "gemini-2.0-flash-exp" - - config = { - } - model = 'gemini-2.0-flash-exp' - if formatting=="raw": - response = client.models.generate_content( - model=model, - contents=messages - ) - + if formatting == "raw": + client = genai.GenerativeModel(model_name=model_name) + response = client.generate_content(contents=messages) + return response.text else: - model = genai.GenerativeModel( - model_name=model, - generation_config=config, - system_instruction=messages[0]["content"], - tools=self._clean_tools_format(tools) + if tools: + client = genai.GenerativeModel( + model_name=model_name, + generation_config=config, + system_instruction=messages[0]["content"], + tools=self._clean_tools_format(tools), ) - chat_session = model.start_chat( - history=self._clean_messages_google(messages)[:-1] - ) - response = chat_session.send_message( - self._clean_messages_google(messages)[-1] - ) - logging.info(response) - return response.text + chat_session = gen_model.start_chat( + history=self._clean_messages_google(messages)[:-1] + ) + response = chat_session.send_message( + self._clean_messages_google(messages)[-1] + ) + return response + else: + gen_model = genai.GenerativeModel( + model_name=model_name, + generation_config=config, + system_instruction=messages[0]["content"], + ) + chat_session = gen_model.start_chat( + history=self._clean_messages_google(messages)[:-1] + ) + response = chat_session.send_message( + self._clean_messages_google(messages)[-1] + ) + return response.text def _raw_gen_stream( - self, - baseself, - model, - messages, - stream=True, - tools=None, - **kwargs - ): - import google.generativeai as genai - genai.configure(api_key=self.api_key) - config = { - } - model = genai.GenerativeModel( - model_name=model, + self, baseself, model, messages, stream=True, tools=None, **kwargs + ): + config = {} + model_name = "gemini-2.0-flash-exp" + + gen_model = genai.GenerativeModel( + model_name=model_name, generation_config=config, - system_instruction=messages[0]["content"] - ) - chat_session = model.start_chat( + system_instruction=messages[0]["content"], + tools=self._clean_tools_format(tools), + ) + chat_session = gen_model.start_chat( history=self._clean_messages_google(messages)[:-1], ) response = chat_session.send_message( - self._clean_messages_google(messages)[-1] - , stream=stream + self._clean_messages_google(messages)[-1], stream=stream ) - for line in response: - if line.text is not None: - yield line.text - + for chunk in response: + if chunk.text is not None: + yield chunk.text + def _supports_tools(self): - return True \ No newline at end of file + return True diff --git a/application/tools/agent.py b/application/tools/agent.py index f4b37d9bf..bbf6bcac3 100644 --- a/application/tools/agent.py +++ b/application/tools/agent.py @@ -89,7 +89,7 @@ def _simple_tool_agent(self, messages): if isinstance(resp, str): yield resp return - if resp.message.content: + if hasattr(resp, "message") and hasattr(resp.message, "content"): yield resp.message.content return @@ -98,7 +98,7 @@ def _simple_tool_agent(self, messages): # If no tool calls are needed, generate the final response if isinstance(resp, str): yield resp - elif resp.message.content: + elif hasattr(resp, "message") and hasattr(resp.message, "content"): yield resp.message.content else: completion = self.llm.gen_stream( diff --git a/application/tools/llm_handler.py b/application/tools/llm_handler.py index 58fce56e4..6be89ad79 100644 --- a/application/tools/llm_handler.py +++ b/application/tools/llm_handler.py @@ -47,23 +47,43 @@ def handle_response(self, agent, resp, tools_dict, messages): class GoogleLLMHandler(LLMHandler): def handle_response(self, agent, resp, tools_dict, messages): - from google.genai import types + import google.generativeai as genai - while resp.content.parts[0].function_call: - function_call_part = resp.candidates[0].content.parts[0] - tool_response, call_id = agent._execute_tool_action( - tools_dict, function_call_part.function_call - ) - function_response_part = types.Part.from_function_response( - name=function_call_part.function_call.name, response=tool_response - ) - - messages.append(function_call_part, function_response_part) + while ( + hasattr(resp.candidates[0].content.parts[0], "function_call") + and resp.candidates[0].content.parts[0].function_call + ): + responses = {} + for part in resp.candidates[0].content.parts: + if hasattr(part, "function_call") and part.function_call: + function_call_part = part + messages.append( + genai.protos.Part( + function_call=genai.protos.FunctionCall( + name=function_call_part.function_call.name, + args=function_call_part.function_call.args, + ) + ) + ) + tool_response, call_id = agent._execute_tool_action( + tools_dict, function_call_part.function_call + ) + responses[function_call_part.function_call.name] = tool_response + response_parts = [ + genai.protos.Part( + function_response=genai.protos.FunctionResponse( + name=tool_name, response={"result": response} + ) + ) + for tool_name, response in responses.items() + ] + if response_parts: + messages.append(response_parts) resp = agent.llm.gen( model=agent.gpt_model, messages=messages, tools=agent.tools ) - return resp + return resp.text def get_llm_handler(llm_type): diff --git a/application/tools/tool_action_parser.py b/application/tools/tool_action_parser.py index b708992a5..254c13b4c 100644 --- a/application/tools/tool_action_parser.py +++ b/application/tools/tool_action_parser.py @@ -1,5 +1,7 @@ import json +from google.protobuf.json_format import MessageToDict + class ToolActionParser: def __init__(self, llm_type): @@ -20,7 +22,8 @@ def _parse_openai_llm(self, call): return tool_id, action_name, call_args def _parse_google_llm(self, call): - call_args = json.loads(call.args) - tool_id = call.name.split("_")[-1] - action_name = call.name.rsplit("_", 1)[0] + call = MessageToDict(call._pb) + call_args = call["args"] + tool_id = call["name"].split("_")[-1] + action_name = call["name"].rsplit("_", 1)[0] return tool_id, action_name, call_args From ec270a3b544a359def3356ec92a77bef8e40cb19 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Fri, 17 Jan 2025 09:51:11 +0530 Subject: [PATCH 05/14] update: requirements --- application/requirements.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/application/requirements.txt b/application/requirements.txt index 08990ab2c..2cc7b1a38 100644 --- a/application/requirements.txt +++ b/application/requirements.txt @@ -14,6 +14,8 @@ esutils==1.0.1 Flask==3.0.3 faiss-cpu==1.9.0.post1 flask-restx==1.3.0 +google-genai==0.5.0 +google-generativeai==0.8.3 gTTS==2.3.2 gunicorn==23.0.0 html2text==2024.2.26 From 904b0bf2da0e215daf4e5bbbbd42a2a30ba89ea1 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Sat, 18 Jan 2025 19:56:25 +0530 Subject: [PATCH 06/14] fix: GoogleLLM, agent and handler according to the new genai SDK --- application/llm/google_ai.py | 199 ++++++++++-------------- application/tools/agent.py | 2 - application/tools/llm_handler.py | 64 ++++---- application/tools/tool_action_parser.py | 9 +- 4 files changed, 116 insertions(+), 158 deletions(-) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 24043d9cd..c14c89b99 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -1,86 +1,61 @@ -import google.generativeai as genai +from google import genai +from google.genai import types from application.core.settings import settings from application.llm.base import BaseLLM class GoogleLLM(BaseLLM): - def __init__(self, api_key=None, user_api_key=None, *args, **kwargs): super().__init__(*args, **kwargs) - self.api_key = settings.API_KEY - genai.configure(api_key=self.api_key) + self.client = genai.Client(api_key="AIzaSyDmbZX65qlQKXcvfMBkJV2KwH82_0yIMlE") def _clean_messages_google(self, messages): cleaned_messages = [] - for message in messages[1:]: - cleaned_messages.append( - { - "role": "model" if message["role"] == "system" else message["role"], - "parts": [message["content"]], - } - ) + for message in messages: + role = message.get("role") + content = message.get("content") + + if role and content is not None: + if isinstance(content, str): + parts = [types.Part.from_text(content)] + elif isinstance(content, list): + parts = content + else: + raise ValueError(f"Unexpected content type: {type(content)}") + + cleaned_messages.append(types.Content(role=role, parts=parts)) + return cleaned_messages - def _clean_tools_format(self, tools_data): - if isinstance(tools_data, list): - return [self._clean_tools_format(item) for item in tools_data] - elif isinstance(tools_data, dict): - if ( - "function" in tools_data - and "type" in tools_data - and tools_data["type"] == "function" - ): - # Handle the case where tools are nested under 'function' - cleaned_function = self._clean_tools_format(tools_data["function"]) - return {"function_declarations": [cleaned_function]} - elif ( - "function" in tools_data - and "type_" in tools_data - and tools_data["type_"] == "function" - ): - # Handle the case where tools are nested under 'function' and type is already 'type_' - cleaned_function = self._clean_tools_format(tools_data["function"]) - return {"function_declarations": [cleaned_function]} - else: - new_tools_data = {} - for key, value in tools_data.items(): - if key == "type": - if value == "string": - new_tools_data["type_"] = "STRING" - elif value == "object": - new_tools_data["type_"] = "OBJECT" - elif key == "additionalProperties": - continue - elif key == "properties": - if isinstance(value, dict): - new_properties = {} - for prop_name, prop_value in value.items(): - if ( - isinstance(prop_value, dict) - and "type" in prop_value - ): - if prop_value["type"] == "string": - new_properties[prop_name] = { - "type_": "STRING", - "description": prop_value.get( - "description" - ), - } - # Add more type mappings as needed - else: - new_properties[prop_name] = ( - self._clean_tools_format(prop_value) - ) - new_tools_data[key] = new_properties - else: - new_tools_data[key] = self._clean_tools_format(value) + def _clean_tools_format(self, tools_list): + genai_tools = [] + for tool_data in tools_list: + if tool_data["type"] == "function": + function = tool_data["function"] + genai_function = dict( + name=function["name"], + description=function["description"], + parameters={ + "type": "OBJECT", + "properties": { + k: { + **v, + "type": v["type"].upper() if v["type"] else None, + } + for k, v in function["parameters"]["properties"].items() + }, + "required": ( + function["parameters"]["required"] + if "required" in function["parameters"] + else [] + ), + }, + ) + genai_tool = types.Tool(function_declarations=[genai_function]) + genai_tools.append(genai_tool) - else: - new_tools_data[key] = self._clean_tools_format(value) - return new_tools_data - else: - return tools_data + return genai_tools def _raw_gen( self, @@ -90,61 +65,51 @@ def _raw_gen( stream=False, tools=None, formatting="openai", - **kwargs + **kwargs, ): - config = {} - model_name = "gemini-2.0-flash-exp" + client = self.client + if formatting == "openai": + messages = self._clean_messages_google(messages) + config = types.GenerateContentConfig() - if formatting == "raw": - client = genai.GenerativeModel(model_name=model_name) - response = client.generate_content(contents=messages) - return response.text + if tools: + cleaned_tools = self._clean_tools_format(tools) + config.tools = cleaned_tools + response = client.models.generate_content( + model=model, + contents=messages, + config=config, + ) + return response else: - if tools: - client = genai.GenerativeModel( - model_name=model_name, - generation_config=config, - system_instruction=messages[0]["content"], - tools=self._clean_tools_format(tools), - ) - chat_session = gen_model.start_chat( - history=self._clean_messages_google(messages)[:-1] - ) - response = chat_session.send_message( - self._clean_messages_google(messages)[-1] - ) - return response - else: - gen_model = genai.GenerativeModel( - model_name=model_name, - generation_config=config, - system_instruction=messages[0]["content"], - ) - chat_session = gen_model.start_chat( - history=self._clean_messages_google(messages)[:-1] - ) - response = chat_session.send_message( - self._clean_messages_google(messages)[-1] - ) - return response.text + response = client.models.generate_content( + model=model, contents=messages, config=config + ) + return response.text def _raw_gen_stream( - self, baseself, model, messages, stream=True, tools=None, **kwargs + self, + baseself, + model, + messages, + stream=True, + tools=None, + formatting="openai", + **kwargs, ): - config = {} - model_name = "gemini-2.0-flash-exp" + client = self.client + if formatting == "openai": + cleaned_messages = self._clean_messages_google(messages) + config = types.GenerateContentConfig() - gen_model = genai.GenerativeModel( - model_name=model_name, - generation_config=config, - system_instruction=messages[0]["content"], - tools=self._clean_tools_format(tools), - ) - chat_session = gen_model.start_chat( - history=self._clean_messages_google(messages)[:-1], - ) - response = chat_session.send_message( - self._clean_messages_google(messages)[-1], stream=stream + if tools: + cleaned_tools = self._clean_tools_format(tools) + config.tools = cleaned_tools + + response = client.models.generate_content_stream( + model=model, + contents=cleaned_messages, + config=config, ) for chunk in response: if chunk.text is not None: diff --git a/application/tools/agent.py b/application/tools/agent.py index bbf6bcac3..209184d2c 100644 --- a/application/tools/agent.py +++ b/application/tools/agent.py @@ -95,7 +95,6 @@ def _simple_tool_agent(self, messages): resp = self.llm_handler.handle_response(self, resp, tools_dict, messages) - # If no tool calls are needed, generate the final response if isinstance(resp, str): yield resp elif hasattr(resp, "message") and hasattr(resp.message, "content"): @@ -110,7 +109,6 @@ def _simple_tool_agent(self, messages): return def gen(self, messages): - # Generate initial response from the LLM if self.llm.supports_tools(): resp = self._simple_tool_agent(messages) for line in resp: diff --git a/application/tools/llm_handler.py b/application/tools/llm_handler.py index 6be89ad79..2383d3f55 100644 --- a/application/tools/llm_handler.py +++ b/application/tools/llm_handler.py @@ -47,43 +47,41 @@ def handle_response(self, agent, resp, tools_dict, messages): class GoogleLLMHandler(LLMHandler): def handle_response(self, agent, resp, tools_dict, messages): - import google.generativeai as genai + from google.genai import types - while ( - hasattr(resp.candidates[0].content.parts[0], "function_call") - and resp.candidates[0].content.parts[0].function_call - ): - responses = {} - for part in resp.candidates[0].content.parts: - if hasattr(part, "function_call") and part.function_call: - function_call_part = part - messages.append( - genai.protos.Part( - function_call=genai.protos.FunctionCall( - name=function_call_part.function_call.name, - args=function_call_part.function_call.args, - ) - ) - ) - tool_response, call_id = agent._execute_tool_action( - tools_dict, function_call_part.function_call - ) - responses[function_call_part.function_call.name] = tool_response - response_parts = [ - genai.protos.Part( - function_response=genai.protos.FunctionResponse( - name=tool_name, response={"result": response} - ) - ) - for tool_name, response in responses.items() - ] - if response_parts: - messages.append(response_parts) - resp = agent.llm.gen( + while True: + response = agent.llm.gen( model=agent.gpt_model, messages=messages, tools=agent.tools ) + if response.candidates and response.candidates[0].content.parts: + tool_call_found = False + for part in response.candidates[0].content.parts: + if part.function_call: + tool_call_found = True + tool_response, call_id = agent._execute_tool_action( + tools_dict, part.function_call + ) + + function_response_part = types.Part.from_function_response( + name=part.function_call.name, + response={"result": tool_response}, + ) + messages.append({"role": "model", "content": [part]}) + messages.append( + {"role": "tool", "content": [function_response_part]} + ) + + if ( + not tool_call_found + and response.candidates[0].content.parts + and response.candidates[0].content.parts[0].text + ): + return response.candidates[0].content.parts[0].text + elif not tool_call_found: + return response.candidates[0].content.parts - return resp.text + else: + return response def get_llm_handler(llm_type): diff --git a/application/tools/tool_action_parser.py b/application/tools/tool_action_parser.py index 254c13b4c..ac0a70c16 100644 --- a/application/tools/tool_action_parser.py +++ b/application/tools/tool_action_parser.py @@ -1,7 +1,5 @@ import json -from google.protobuf.json_format import MessageToDict - class ToolActionParser: def __init__(self, llm_type): @@ -22,8 +20,7 @@ def _parse_openai_llm(self, call): return tool_id, action_name, call_args def _parse_google_llm(self, call): - call = MessageToDict(call._pb) - call_args = call["args"] - tool_id = call["name"].split("_")[-1] - action_name = call["name"].rsplit("_", 1)[0] + call_args = call.args + tool_id = call.name.split("_")[-1] + action_name = call.name.rsplit("_", 1)[0] return tool_id, action_name, call_args From 4a331db5fc110d3ff64e6a9966b49f51e2789414 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Sat, 18 Jan 2025 20:00:51 +0530 Subject: [PATCH 07/14] fix: api_key attribute --- application/llm/google_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index c14c89b99..e9188987c 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -8,7 +8,7 @@ class GoogleLLM(BaseLLM): def __init__(self, api_key=None, user_api_key=None, *args, **kwargs): super().__init__(*args, **kwargs) - self.client = genai.Client(api_key="AIzaSyDmbZX65qlQKXcvfMBkJV2KwH82_0yIMlE") + self.client = genai.Client(api_key=api_key) def _clean_messages_google(self, messages): cleaned_messages = [] From a741388447dfc6763d435639175e90cfb0b2b334 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 19 Jan 2025 22:21:50 +0000 Subject: [PATCH 08/14] fix: system message --- application/llm/google_ai.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index e9188987c..9ebc0a5e3 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -71,6 +71,9 @@ def _raw_gen( if formatting == "openai": messages = self._clean_messages_google(messages) config = types.GenerateContentConfig() + if messages[0].role == "system": + config.system_instruction = messages[0].parts[0].text + messages = messages[1:] if tools: cleaned_tools = self._clean_tools_format(tools) @@ -101,6 +104,9 @@ def _raw_gen_stream( if formatting == "openai": cleaned_messages = self._clean_messages_google(messages) config = types.GenerateContentConfig() + if messages[0].role == "system": + config.system_instruction = messages[0].parts[0].text + messages = messages[1:] if tools: cleaned_tools = self._clean_tools_format(tools) From d441d5763fd6e265abda203dc2d281a6439c2ce7 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Mon, 20 Jan 2025 19:44:14 +0530 Subject: [PATCH 09/14] fix: decorators + client error --- application/cache.py | 2 +- application/llm/base.py | 35 +++++++++++++++++++++++++------- application/llm/google_ai.py | 34 +++++++++++++++++++++++++------ application/tools/llm_handler.py | 11 +++++++--- 4 files changed, 65 insertions(+), 17 deletions(-) diff --git a/application/cache.py b/application/cache.py index 76b594c93..72bab2b93 100644 --- a/application/cache.py +++ b/application/cache.py @@ -68,7 +68,7 @@ def wrapper(self, model, messages, stream, tools=None, *args, **kwargs): def stream_cache(func): - def wrapper(self, model, messages, stream, *args, **kwargs): + def wrapper(self, model, messages, stream, tools=None, *args, **kwargs): cache_key = gen_cache_key(messages) logger.info(f"Stream cache key: {cache_key}") diff --git a/application/llm/base.py b/application/llm/base.py index b9b0e5243..e687e567b 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod + +from application.cache import gen_cache, stream_cache from application.usage import gen_token_usage, stream_token_usage -from application.cache import stream_cache, gen_cache class BaseLLM(ABC): @@ -18,18 +19,38 @@ def _raw_gen(self, model, messages, stream, tools, *args, **kwargs): def gen(self, model, messages, stream=False, tools=None, *args, **kwargs): decorators = [gen_token_usage, gen_cache] - return self._apply_decorator(self._raw_gen, decorators=decorators, model=model, messages=messages, stream=stream, tools=tools, *args, **kwargs) + return self._apply_decorator( + self._raw_gen, + decorators=decorators, + model=model, + messages=messages, + stream=stream, + tools=tools, + *args, + **kwargs + ) @abstractmethod def _raw_gen_stream(self, model, messages, stream, *args, **kwargs): pass - def gen_stream(self, model, messages, stream=True, *args, **kwargs): + def gen_stream(self, model, messages, stream=True, tools=None, *args, **kwargs): decorators = [stream_cache, stream_token_usage] - return self._apply_decorator(self._raw_gen_stream, decorators=decorators, model=model, messages=messages, stream=stream, *args, **kwargs) - + return self._apply_decorator( + self._raw_gen_stream, + decorators=decorators, + model=model, + messages=messages, + stream=stream, + tools=tools, + *args, + **kwargs + ) + def supports_tools(self): - return hasattr(self, '_supports_tools') and callable(getattr(self, '_supports_tools')) + return hasattr(self, "_supports_tools") and callable( + getattr(self, "_supports_tools") + ) def _supports_tools(self): - raise NotImplementedError("Subclass must implement _supports_tools method") \ No newline at end of file + raise NotImplementedError("Subclass must implement _supports_tools method") diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 9ebc0a5e3..7e67a4cdd 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -8,7 +8,8 @@ class GoogleLLM(BaseLLM): def __init__(self, api_key=None, user_api_key=None, *args, **kwargs): super().__init__(*args, **kwargs) - self.client = genai.Client(api_key=api_key) + self.api_key = api_key + self.user_api_key = user_api_key def _clean_messages_google(self, messages): cleaned_messages = [] @@ -16,11 +17,32 @@ def _clean_messages_google(self, messages): role = message.get("role") content = message.get("content") + parts = [] if role and content is not None: if isinstance(content, str): parts = [types.Part.from_text(content)] elif isinstance(content, list): - parts = content + for item in content: + if "text" in item: + parts.append(types.Part.from_text(item["text"])) + elif "function_call" in item: + parts.append( + types.Part.from_function_call( + name=item["function_call"]["name"], + args=item["function_call"]["args"], + ) + ) + elif "function_response" in item: + parts.append( + types.Part.from_function_response( + name=item["function_response"]["name"], + response=item["function_response"]["response"], + ) + ) + else: + raise ValueError( + f"Unexpected content dictionary format:{item}" + ) else: raise ValueError(f"Unexpected content type: {type(content)}") @@ -67,7 +89,7 @@ def _raw_gen( formatting="openai", **kwargs, ): - client = self.client + client = genai.Client(api_key=self.api_key) if formatting == "openai": messages = self._clean_messages_google(messages) config = types.GenerateContentConfig() @@ -100,9 +122,9 @@ def _raw_gen_stream( formatting="openai", **kwargs, ): - client = self.client + client = genai.Client(api_key=self.api_key) if formatting == "openai": - cleaned_messages = self._clean_messages_google(messages) + messages = self._clean_messages_google(messages) config = types.GenerateContentConfig() if messages[0].role == "system": config.system_instruction = messages[0].parts[0].text @@ -114,7 +136,7 @@ def _raw_gen_stream( response = client.models.generate_content_stream( model=model, - contents=cleaned_messages, + contents=messages, config=config, ) for chunk in response: diff --git a/application/tools/llm_handler.py b/application/tools/llm_handler.py index 2383d3f55..cc7494c02 100644 --- a/application/tools/llm_handler.py +++ b/application/tools/llm_handler.py @@ -61,14 +61,19 @@ def handle_response(self, agent, resp, tools_dict, messages): tool_response, call_id = agent._execute_tool_action( tools_dict, part.function_call ) - function_response_part = types.Part.from_function_response( name=part.function_call.name, response={"result": tool_response}, ) - messages.append({"role": "model", "content": [part]}) + + messages.append( + {"role": "model", "content": [part.to_json_dict()]} + ) messages.append( - {"role": "tool", "content": [function_response_part]} + { + "role": "tool", + "content": [function_response_part.to_json_dict()], + } ) if ( From 1086bfe1bac1a13b3232f53d3858c944e62328c5 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Tue, 21 Jan 2025 07:19:02 +0530 Subject: [PATCH 10/14] fix: wrong role in req messages --- application/llm/google_ai.py | 3 +++ application/retriever/classic_rag.py | 13 +++++-------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 7e67a4cdd..5b1eeb740 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -17,6 +17,9 @@ def _clean_messages_google(self, messages): role = message.get("role") content = message.get("content") + if role == "assistant": + role = "model" + parts = [] if role and content is not None: if isinstance(content, str): diff --git a/application/retriever/classic_rag.py b/application/retriever/classic_rag.py index 2e3555137..b3735a963 100644 --- a/application/retriever/classic_rag.py +++ b/application/retriever/classic_rag.py @@ -5,7 +5,6 @@ from application.vectorstore.vector_creator import VectorCreator - class ClassicRAG(BaseRetriever): def __init__( @@ -74,13 +73,11 @@ def gen(self): if len(self.chat_history) > 1: for i in self.chat_history: - if "prompt" in i and "response" in i: - messages_combine.append( - {"role": "user", "content": i["prompt"]} - ) - messages_combine.append( - {"role": "system", "content": i["response"]} - ) + if "prompt" in i and "response" in i: + messages_combine.append({"role": "user", "content": i["prompt"]}) + messages_combine.append( + {"role": "assistant", "content": i["response"]} + ) messages_combine.append({"role": "user", "content": self.question}) # llm = LLMCreator.create_llm( # settings.LLM_NAME, api_key=settings.API_KEY, user_api_key=self.user_api_key From 94f682e461fc49efbd1bddb07ee553835268dc8e Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 21 Jan 2025 09:31:38 +0000 Subject: [PATCH 11/14] fix roles in retriever layer --- application/retriever/brave_search.py | 10 ++++------ application/retriever/duckduck_search.py | 6 ++---- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/application/retriever/brave_search.py b/application/retriever/brave_search.py index 3d9ae89e6..efcae8ab9 100644 --- a/application/retriever/brave_search.py +++ b/application/retriever/brave_search.py @@ -74,12 +74,10 @@ def gen(self): if len(self.chat_history) > 1: for i in self.chat_history: if "prompt" in i and "response" in i: - messages_combine.append( - {"role": "user", "content": i["prompt"]} - ) - messages_combine.append( - {"role": "system", "content": i["response"]} - ) + messages_combine.append({"role": "user", "content": i["prompt"]}) + messages_combine.append( + {"role": "assistant", "content": i["response"]} + ) messages_combine.append({"role": "user", "content": self.question}) llm = LLMCreator.create_llm( diff --git a/application/retriever/duckduck_search.py b/application/retriever/duckduck_search.py index fa19ead03..321c6fd96 100644 --- a/application/retriever/duckduck_search.py +++ b/application/retriever/duckduck_search.py @@ -91,11 +91,9 @@ def gen(self): if len(self.chat_history) > 1: for i in self.chat_history: if "prompt" in i and "response" in i: + messages_combine.append({"role": "user", "content": i["prompt"]}) messages_combine.append( - {"role": "user", "content": i["prompt"]} - ) - messages_combine.append( - {"role": "system", "content": i["response"]} + {"role": "assistant", "content": i["response"]} ) messages_combine.append({"role": "user", "content": self.question}) From c0c60a487589588ac28117923f98f569783c1ed9 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 21 Jan 2025 09:37:11 +0000 Subject: [PATCH 12/14] fix: ruff linting --- application/llm/google_ai.py | 1 - 1 file changed, 1 deletion(-) diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index 5b1eeb740..ae8880421 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -1,7 +1,6 @@ from google import genai from google.genai import types -from application.core.settings import settings from application.llm.base import BaseLLM From 9f073fcbcf1e50107bfa0b0a5eab0d929670bfea Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 21 Jan 2025 09:39:02 +0000 Subject: [PATCH 13/14] fix: tests --- application/cache.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/application/cache.py b/application/cache.py index 72bab2b93..0b5d72da6 100644 --- a/application/cache.py +++ b/application/cache.py @@ -69,7 +69,7 @@ def wrapper(self, model, messages, stream, tools=None, *args, **kwargs): def stream_cache(func): def wrapper(self, model, messages, stream, tools=None, *args, **kwargs): - cache_key = gen_cache_key(messages) + cache_key = gen_cache_key(messages, model, tools) logger.info(f"Stream cache key: {cache_key}") redis_client = get_redis_instance() @@ -86,7 +86,7 @@ def wrapper(self, model, messages, stream, tools=None, *args, **kwargs): except redis.ConnectionError as e: logger.error(f"Redis connection error: {e}") - result = func(self, model, messages, stream, *args, **kwargs) + result = func(self, model, messages, stream, tools=tools, *args, **kwargs) stream_cache_data = [] for chunk in result: From 43340c4aa8b56b623846aabea250709e232d582e Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 21 Jan 2025 09:44:53 +0000 Subject: [PATCH 14/14] fix: finale test --- application/cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/application/cache.py b/application/cache.py index 0b5d72da6..80dee4f48 100644 --- a/application/cache.py +++ b/application/cache.py @@ -33,7 +33,7 @@ def gen_cache_key(messages, model="docgpt", tools=None): if not all(isinstance(msg, dict) for msg in messages): raise ValueError("All messages must be dictionaries.") messages_str = json.dumps(messages) - tools_str = json.dumps(tools) if tools else "" + tools_str = json.dumps(str(tools)) if tools else "" combined = f"{model}_{messages_str}_{tools_str}" cache_key = get_hash(combined) return cache_key