LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2026-02-26 22:31:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: GoogleLLM, agent and handler according to the new genai SDK

Browse Source
This commit is contained in:
Siddhant Rai
2025-01-18 19:56:25 +05:30
parent ec270a3b54
commit 904b0bf2da
4 changed files with 116 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files

199
application/llm/google_ai.py
View File

@@ -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.core.settings import settings
from application.llm.base import BaseLLM from application.llm.base import BaseLLM
class GoogleLLM(BaseLLM): class GoogleLLM(BaseLLM):
def __init__(self, api_key=None, user_api_key=None, *args, **kwargs): def __init__(self, api_key=None, user_api_key=None, *args, **kwargs):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.api_key = settings.API_KEY self.client = genai.Client(api_key="AIzaSyDmbZX65qlQKXcvfMBkJV2KwH82_0yIMlE")
genai.configure(api_key=self.api_key)
def _clean_messages_google(self, messages): def _clean_messages_google(self, messages):
cleaned_messages = [] cleaned_messages = []
for message in messages[1:]: for message in messages:
cleaned_messages.append( role = message.get("role")
{ content = message.get("content")
"role": "model" if message["role"] == "system" else message["role"],
"parts": [message["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 return cleaned_messages
def _clean_tools_format(self, tools_data): def _clean_tools_format(self, tools_list):
if isinstance(tools_data, list): genai_tools = []
return [self._clean_tools_format(item) for item in tools_data] for tool_data in tools_list:
elif isinstance(tools_data, dict): if tool_data["type"] == "function":
if ( function = tool_data["function"]
"function" in tools_data genai_function = dict(
and "type" in tools_data name=function["name"],
and tools_data["type"] == "function" description=function["description"],
): parameters={
# Handle the case where tools are nested under 'function' "type": "OBJECT",
cleaned_function = self._clean_tools_format(tools_data["function"]) "properties": {
return {"function_declarations": [cleaned_function]} k: {
elif ( **v,
"function" in tools_data "type": v["type"].upper() if v["type"] else None,
and "type_" in tools_data }
and tools_data["type_"] == "function" for k, v in function["parameters"]["properties"].items()
): },
# Handle the case where tools are nested under 'function' and type is already 'type_' "required": (
cleaned_function = self._clean_tools_format(tools_data["function"]) function["parameters"]["required"]
return {"function_declarations": [cleaned_function]} if "required" in function["parameters"]
else: else []
new_tools_data = {} ),
for key, value in tools_data.items(): },
if key == "type": )
if value == "string": genai_tool = types.Tool(function_declarations=[genai_function])
new_tools_data["type_"] = "STRING" genai_tools.append(genai_tool)
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)
else: return genai_tools
new_tools_data[key] = self._clean_tools_format(value)
return new_tools_data
else:
return tools_data
def _raw_gen( def _raw_gen(
self, self,
@@ -90,61 +65,51 @@ class GoogleLLM(BaseLLM):
stream=False, stream=False,
tools=None, tools=None,
formatting="openai", formatting="openai",
**kwargs **kwargs,
): ):
config = {} client = self.client
model_name = "gemini-2.0-flash-exp" if formatting == "openai":
messages = self._clean_messages_google(messages)
config = types.GenerateContentConfig()
if formatting == "raw": if tools:
client = genai.GenerativeModel(model_name=model_name) cleaned_tools = self._clean_tools_format(tools)
response = client.generate_content(contents=messages) config.tools = cleaned_tools
return response.text response = client.models.generate_content(
model=model,
contents=messages,
config=config,
)
return response
else: else:
if tools: response = client.models.generate_content(
client = genai.GenerativeModel( model=model, contents=messages, config=config
model_name=model_name, )
generation_config=config, return response.text
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
def _raw_gen_stream( 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 = {} client = self.client
model_name = "gemini-2.0-flash-exp" if formatting == "openai":
cleaned_messages = self._clean_messages_google(messages)
config = types.GenerateContentConfig()
gen_model = genai.GenerativeModel( if tools:
model_name=model_name, cleaned_tools = self._clean_tools_format(tools)
generation_config=config, config.tools = cleaned_tools
system_instruction=messages[0]["content"],
tools=self._clean_tools_format(tools), response = client.models.generate_content_stream(
) model=model,
chat_session = gen_model.start_chat( contents=cleaned_messages,
history=self._clean_messages_google(messages)[:-1], config=config,
)
response = chat_session.send_message(
self._clean_messages_google(messages)[-1], stream=stream
) )
for chunk in response: for chunk in response:
if chunk.text is not None: if chunk.text is not None:

2
application/tools/agent.py
View File

@@ -95,7 +95,6 @@ class Agent:
resp = self.llm_handler.handle_response(self, resp, tools_dict, 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): if isinstance(resp, str):
yield resp yield resp
elif hasattr(resp, "message") and hasattr(resp.message, "content"): elif hasattr(resp, "message") and hasattr(resp.message, "content"):
@@ -110,7 +109,6 @@ class Agent:
return return
def gen(self, messages): def gen(self, messages):
# Generate initial response from the LLM
if self.llm.supports_tools(): if self.llm.supports_tools():
resp = self._simple_tool_agent(messages) resp = self._simple_tool_agent(messages)
for line in resp: for line in resp:

64
application/tools/llm_handler.py
View File

@@ -47,43 +47,41 @@ class OpenAILLMHandler(LLMHandler):
class GoogleLLMHandler(LLMHandler): class GoogleLLMHandler(LLMHandler):
def handle_response(self, agent, resp, tools_dict, messages): def handle_response(self, agent, resp, tools_dict, messages):
import google.generativeai as genai from google.genai import types
while ( while True:
hasattr(resp.candidates[0].content.parts[0], "function_call") response = agent.llm.gen(
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 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
)
return resp.text 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
else:
return response
def get_llm_handler(llm_type): def get_llm_handler(llm_type):

9
application/tools/tool_action_parser.py
View File

@@ -1,7 +1,5 @@
import json import json
from google.protobuf.json_format import MessageToDict
class ToolActionParser: class ToolActionParser:
def __init__(self, llm_type): def __init__(self, llm_type):
@@ -22,8 +20,7 @@ class ToolActionParser:
return tool_id, action_name, call_args return tool_id, action_name, call_args
def _parse_google_llm(self, call): def _parse_google_llm(self, call):
call = MessageToDict(call._pb) call_args = call.args
call_args = call["args"] tool_id = call.name.split("_")[-1]
tool_id = call["name"].split("_")[-1] action_name = call.name.rsplit("_", 1)[0]
action_name = call["name"].rsplit("_", 1)[0]
return tool_id, action_name, call_args return tool_id, action_name, call_args