fix: google parser, llm handler and other errors

This commit is contained in:
Siddhant Rai
2025-01-17 09:22:41 +05:30
parent 811dfecf98
commit c97d1e3363
4 changed files with 141 additions and 102 deletions

View File

@@ -89,7 +89,7 @@ class Agent:
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 @@ class Agent:
# 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(

View File

@@ -47,23 +47,43 @@ class OpenAILLMHandler(LLMHandler):
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):

View File

@@ -1,5 +1,7 @@
import json
from google.protobuf.json_format import MessageToDict
class ToolActionParser:
def __init__(self, llm_type):
@@ -20,7 +22,8 @@ class ToolActionParser:
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