my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from typing import Dict, List, Optional
#3
#4	try:
#5	from google import genai
#6	from google.genai import types
#7	except ImportError:
#8	raise ImportError("The 'google-genai' library is required. Please install it using 'pip install google-genai'.")
#9
#10	from mem0.configs.llms.base import BaseLlmConfig
#11	from mem0.llms.base import LLMBase
#12
#13
#14	class GeminiLLM(LLMBase):
#15	def __init__(self, config: Optional[BaseLlmConfig] = None):
#16	super().__init__(config)
#17
#18	if not self.config.model:
#19	self.config.model = "gemini-2.0-flash"
#20
#21	api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
#22	self.client = genai.Client(api_key=api_key)
#23
#24	def _parse_response(self, response, tools):
#25	"""
#26	Process the response based on whether tools are used or not.
#27
#28	Args:
#29	response: The raw response from API.
#30	tools: The list of tools provided in the request.
#31
#32	Returns:
#33	str or dict: The processed response.
#34	"""
#35	if tools:
#36	processed_response = {
#37	"content": None,
#38	"tool_calls": [],
#39	}
#40
#41	# Extract content from the first candidate
#42	if response.candidates and response.candidates[0].content.parts:
#43	for part in response.candidates[0].content.parts:
#44	if hasattr(part, "text") and part.text:
#45	processed_response["content"] = part.text
#46	break
#47
#48	# Extract function calls
#49	if response.candidates and response.candidates[0].content.parts:
#50	for part in response.candidates[0].content.parts:
#51	if hasattr(part, "function_call") and part.function_call:
#52	fn = part.function_call
#53	processed_response["tool_calls"].append(
#54	{
#55	"name": fn.name,
#56	"arguments": dict(fn.args) if fn.args else {},
#57	}
#58	)
#59
#60	return processed_response
#61	else:
#62	if response.candidates and response.candidates[0].content.parts:
#63	for part in response.candidates[0].content.parts:
#64	if hasattr(part, "text") and part.text:
#65	return part.text
#66	return ""
#67
#68	def _reformat_messages(self, messages: List[Dict[str, str]]):
#69	"""
#70	Reformat messages for Gemini.
#71
#72	Args:
#73	messages: The list of messages provided in the request.
#74
#75	Returns:
#76	tuple: (system_instruction, contents_list)
#77	"""
#78	system_instruction = None
#79	contents = []
#80
#81	for message in messages:
#82	if message["role"] == "system":
#83	system_instruction = message["content"]
#84	else:
#85	content = types.Content(
#86	parts=[types.Part(text=message["content"])],
#87	role=message["role"],
#88	)
#89	contents.append(content)
#90
#91	return system_instruction, contents
#92
#93	def _reformat_tools(self, tools: Optional[List[Dict]]):
#94	"""
#95	Reformat tools for Gemini.
#96
#97	Args:
#98	tools: The list of tools provided in the request.
#99
#100	Returns:
#101	list: The list of tools in the required format.
#102	"""
#103
#104	def remove_additional_properties(data):
#105	"""Recursively removes 'additionalProperties' from nested dictionaries."""
#106	if isinstance(data, dict):
#107	filtered_dict = {
#108	key: remove_additional_properties(value)
#109	for key, value in data.items()
#110	if not (key == "additionalProperties")
#111	}
#112	return filtered_dict
#113	else:
#114	return data
#115
#116	if tools:
#117	function_declarations = []
#118	for tool in tools:
#119	func = tool["function"].copy()
#120	cleaned_func = remove_additional_properties(func)
#121
#122	function_declaration = types.FunctionDeclaration(
#123	name=cleaned_func["name"],
#124	description=cleaned_func.get("description", ""),
#125	parameters=cleaned_func.get("parameters", {}),
#126	)
#127	function_declarations.append(function_declaration)
#128
#129	tool_obj = types.Tool(function_declarations=function_declarations)
#130	return [tool_obj]
#131	else:
#132	return None
#133
#134	def generate_response(
#135	self,
#136	messages: List[Dict[str, str]],
#137	response_format=None,
#138	tools: Optional[List[Dict]] = None,
#139	tool_choice: str = "auto",
#140	):
#141	"""
#142	Generate a response based on the given messages using Gemini.
#143
#144	Args:
#145	messages (list): List of message dicts containing 'role' and 'content'.
#146	response_format (str or object, optional): Format for the response. Defaults to "text".
#147	tools (list, optional): List of tools that the model can call. Defaults to None.
#148	tool_choice (str, optional): Tool choice method. Defaults to "auto".
#149
#150	Returns:
#151	str: The generated response.
#152	"""
#153
#154	# Extract system instruction and reformat messages
#155	system_instruction, contents = self._reformat_messages(messages)
#156
#157	# Prepare generation config
#158	config_params = {
#159	"temperature": self.config.temperature,
#160	"max_output_tokens": self.config.max_tokens,
#161	"top_p": self.config.top_p,
#162	}
#163
#164	# Add system instruction to config if present
#165	if system_instruction:
#166	config_params["system_instruction"] = system_instruction
#167
#168	if response_format is not None and response_format["type"] == "json_object":
#169	config_params["response_mime_type"] = "application/json"
#170	if "schema" in response_format:
#171	config_params["response_schema"] = response_format["schema"]
#172
#173	if tools:
#174	formatted_tools = self._reformat_tools(tools)
#175	config_params["tools"] = formatted_tools
#176
#177	if tool_choice:
#178	if tool_choice == "auto":
#179	mode = types.FunctionCallingConfigMode.AUTO
#180	elif tool_choice == "any":
#181	mode = types.FunctionCallingConfigMode.ANY
#182	else:
#183	mode = types.FunctionCallingConfigMode.NONE
#184
#185	tool_config = types.ToolConfig(
#186	function_calling_config=types.FunctionCallingConfig(
#187	mode=mode,
#188	allowed_function_names=(
#189	[tool["function"]["name"] for tool in tools] if tool_choice == "any" else None
#190	),
#191	)
#192	)
#193	config_params["tool_config"] = tool_config
#194
#195	generation_config = types.GenerateContentConfig(**config_params)
#196
#197	response = self.client.models.generate_content(
#198	model=self.config.model, contents=contents, config=generation_config
#199	)
#200
#201	return self._parse_response(response, tools)
#202

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public