fromllama_index.llms.geminiimportGeminillm=Gemini(model="models/gemini-ultra",api_key="YOUR_API_KEY")resp=llm.complete("Write a poem about a magic backpack")print(resp)
@deprecated.deprecated(reason=("Should use `llama-index-llms-google-genai` instead, using Google's latest unified SDK. ""See: https://docs.llamaindex.org.cn/en/stable/examples/llm/google_genai/"))classGemini(FunctionCallingLLM):""" Gemini LLM. Examples: `pip install llama-index-llms-gemini` ```python from llama_index.llms.gemini import Gemini llm = Gemini(model="models/gemini-ultra", api_key="YOUR_API_KEY") resp = llm.complete("Write a poem about a magic backpack") print(resp) ``` """model:str=Field(default=GEMINI_MODELS[0],description="The Gemini model to use.")temperature:float=Field(default=DEFAULT_TEMPERATURE,description="The temperature to use during generation.",ge=0.0,le=2.0,)max_tokens:int=Field(default=DEFAULT_NUM_OUTPUTS,description="The number of tokens to generate.",gt=0,)generate_kwargs:dict=Field(default_factory=dict,description="Kwargs for generation.")_model:genai.GenerativeModel=PrivateAttr()_model_meta:genai.types.Model=PrivateAttr()_request_options:Optional[genai.types.RequestOptions]=PrivateAttr()def__init__(self,api_key:Optional[str]=None,model:str=GEMINI_MODELS[0],temperature:float=DEFAULT_TEMPERATURE,max_tokens:Optional[int]=None,generation_config:Optional[genai.types.GenerationConfigDict]=None,safety_settings:Optional[genai.types.SafetySettingDict]=None,callback_manager:Optional[CallbackManager]=None,api_base:Optional[str]=None,transport:Optional[str]=None,model_name:Optional[str]=None,default_headers:Optional[Dict[str,str]]=None,request_options:Optional[genai.types.RequestOptions]=None,**generate_kwargs:Any,):"""Creates a new Gemini model interface."""ifmodel_nameisnotNone:warnings.warn("model_name is deprecated, please use model instead",DeprecationWarning,)model=model_name# API keys are optional. The API can be authorised via OAuth (detected# environmentally) or by the GOOGLE_API_KEY environment variable.config_params:Dict[str,Any]={"api_key":api_keyoros.getenv("GOOGLE_API_KEY"),}ifapi_base:config_params["client_options"]={"api_endpoint":api_base}iftransport:config_params["transport"]=transportifdefault_headers:default_metadata=[]forkey,valueindefault_headers.items():default_metadata.append((key,value))# `default_metadata` contains (key, value) pairs that will be sent with every request.# When using `transport="rest"`, these will be sent as HTTP headers.config_params["default_metadata"]=default_metadata# transport: A string, one of: [`rest`, `grpc`, `grpc_asyncio`].genai.configure(**config_params)base_gen_config=generation_configifgeneration_configelse{}# Explicitly passed args take precedence over the generation_config.final_gen_config=cast(generation_types.GenerationConfigDict,{"temperature":temperature,**base_gen_config},)model_meta=genai.get_model(model)genai_model=genai.GenerativeModel(model_name=model,generation_config=final_gen_config,safety_settings=safety_settings,)supported_methods=model_meta.supported_generation_methodsif"generateContent"notinsupported_methods:raiseValueError(f"Model {model} does not support content generation, only "f"{supported_methods}.")ifnotmax_tokens:max_tokens=model_meta.output_token_limitelse:max_tokens=min(max_tokens,model_meta.output_token_limit)super().__init__(model=model,temperature=temperature,max_tokens=max_tokens,generate_kwargs=generate_kwargs,callback_manager=callback_manager,)self._model_meta=model_metaself._model=genai_modelself._request_options=request_optionsself._is_function_call_model=is_function_calling_model(model)@classmethoddefclass_name(cls)->str:return"Gemini_LLM"@propertydefmetadata(self)->LLMMetadata:total_tokens=self._model_meta.input_token_limit+self.max_tokensreturnLLMMetadata(context_window=total_tokens,num_output=self.max_tokens,model_name=self.model,is_chat_model=True,# All gemini models support function callingis_function_calling_model=self._is_function_call_model,)@llm_completion_callback()defcomplete(self,prompt:str,formatted:bool=False,**kwargs:Any)->CompletionResponse:request_options=self._request_optionsorkwargs.pop("request_options",None)result=self._model.generate_content(prompt,request_options=request_options,**kwargs)returncompletion_from_gemini_response(result)@llm_completion_callback()asyncdefacomplete(self,prompt:str,formatted:bool=False,**kwargs:Any)->CompletionResponse:request_options=self._request_optionsorkwargs.pop("request_options",None)result=awaitself._model.generate_content_async(prompt,request_options=request_options,**kwargs)returncompletion_from_gemini_response(result)@llm_completion_callback()defstream_complete(self,prompt:str,formatted:bool=False,**kwargs:Any)->CompletionResponseGen:request_options=self._request_optionsorkwargs.pop("request_options",None)defgen():text=""it=self._model.generate_content(prompt,stream=True,request_options=request_options,**kwargs)forrinit:delta=r.textor""text+=deltayieldcompletion_from_gemini_response(r,text=text,delta=delta)returngen()@llm_completion_callback()defastream_complete(self,prompt:str,formatted:bool=False,**kwargs:Any)->CompletionResponseAsyncGen:request_options=self._request_optionsorkwargs.pop("request_options",None)asyncdefgen():text=""it=awaitself._model.generate_content_async(prompt,stream=True,request_options=request_options,**kwargs)asyncforrinit:delta=r.textor""text+=deltayieldcompletion_from_gemini_response(r,text=text,delta=delta)returngen()@llm_chat_callback()defchat(self,messages:Sequence[ChatMessage],**kwargs:Any)->ChatResponse:request_options=self._request_optionsorkwargs.pop("request_options",None)merged_messages=merge_neighboring_same_role_messages(messages)*history,next_msg=map(chat_message_to_gemini,merged_messages)chat=self._model.start_chat(history=history)response=chat.send_message(next_msg,request_options=request_options,**kwargs,)returnchat_from_gemini_response(response)@llm_chat_callback()asyncdefachat(self,messages:Sequence[ChatMessage],**kwargs:Any)->ChatResponse:request_options=self._request_optionsorkwargs.pop("request_options",None)merged_messages=merge_neighboring_same_role_messages(messages)*history,next_msg=map(chat_message_to_gemini,merged_messages)chat=self._model.start_chat(history=history)response=awaitchat.send_message_async(next_msg,request_options=request_options,**kwargs)returnchat_from_gemini_response(response)@llm_chat_callback()defstream_chat(self,messages:Sequence[ChatMessage],**kwargs:Any)->ChatResponseGen:request_options=self._request_optionsorkwargs.pop("request_options",None)merged_messages=merge_neighboring_same_role_messages(messages)*history,next_msg=map(chat_message_to_gemini,merged_messages)chat=self._model.start_chat(history=history)response=chat.send_message(next_msg,stream=True,request_options=request_options,**kwargs)defgen()->ChatResponseGen:content=""existing_tool_calls=[]forrinresponse:top_candidate=r.candidates[0]content_delta=top_candidate.content.parts[0].textcontent+=content_deltallama_resp=chat_from_gemini_response(r)existing_tool_calls.extend(llama_resp.message.additional_kwargs.get("tool_calls",[]))llama_resp.delta=content_deltallama_resp.message.content=contentllama_resp.message.additional_kwargs["tool_calls"]=existing_tool_callsyieldllama_respreturngen()@llm_chat_callback()asyncdefastream_chat(self,messages:Sequence[ChatMessage],**kwargs:Any)->ChatResponseAsyncGen:request_options=self._request_optionsorkwargs.pop("request_options",None)merged_messages=merge_neighboring_same_role_messages(messages)*history,next_msg=map(chat_message_to_gemini,merged_messages)chat=self._model.start_chat(history=history)response=awaitchat.send_message_async(next_msg,stream=True,request_options=request_options,**kwargs)asyncdefgen()->ChatResponseAsyncGen:content=""existing_tool_calls=[]asyncforrinresponse:top_candidate=r.candidates[0]content_delta=top_candidate.content.parts[0].textcontent+=content_deltallama_resp=chat_from_gemini_response(r)existing_tool_calls.extend(llama_resp.message.additional_kwargs.get("tool_calls",[]))llama_resp.delta=content_deltallama_resp.message.content=contentllama_resp.message.additional_kwargs["tool_calls"]=existing_tool_callsyieldllama_respreturngen()def_prepare_chat_with_tools(self,tools:Sequence["BaseTool"],user_msg:Optional[Union[str,ChatMessage]]=None,chat_history:Optional[List[ChatMessage]]=None,verbose:bool=False,allow_parallel_tool_calls:bool=False,tool_choice:Union[str,dict]="auto",strict:Optional[bool]=None,**kwargs:Any,)->Dict[str,Any]:"""Predict and call the tool."""fromgoogle.generativeai.typesimportFunctionDeclaration,ToolDictfromgoogle.generativeai.types.content_typesimportFunctionCallingModeiftool_choice=="auto":tool_mode=FunctionCallingMode.AUTOeliftool_choice=="none":tool_mode=FunctionCallingMode.NONEelse:tool_mode=FunctionCallingMode.ANYtool_config={"function_calling_config":{"mode":tool_mode,}}iftool_choicenotin["auto","none"]:ifisinstance(tool_choice,dict):raiseValueError("Gemini does not support tool_choice as a dict")# assume that the user wants a tool call to be made# if the tool choice is not in the list of tools, then we will make a tool call to all tools# otherwise, we will make a tool call to the tool choicetool_names=[tool.metadata.namefortoolintools]iftool_choicenotintool_names:tool_config["function_calling_config"]["allowed_function_names"]=tool_nameselse:tool_config["function_calling_config"]["allowed_function_names"]=[tool_choice]tool_declarations=[]fortoolintools:descriptions={}forparam_name,param_schemaintool.metadata.get_parameters_dict()["properties"].items():param_description=param_schema.get("description",None)ifparam_description:descriptions[param_name]=param_descriptiontool.metadata.fn_schema.__doc__=tool.metadata.descriptiontool_declarations.append(FunctionDeclaration.from_function(tool.metadata.fn_schema,descriptions))ifisinstance(user_msg,str):user_msg=ChatMessage(role=MessageRole.USER,content=user_msg)messages=chat_historyor[]ifuser_msg:messages.append(user_msg)return{"messages":messages,"tools":(ToolDict(function_declarations=tool_declarations)iftool_declarationselseNone),"tool_config":tool_config,**kwargs,}defget_tool_calls_from_response(self,response:ChatResponse,error_on_no_tool_call:bool=True,**kwargs:Any,)->List[ToolSelection]:"""Predict and call the tool."""tool_calls=response.message.additional_kwargs.get("tool_calls",[])iflen(tool_calls)<1:iferror_on_no_tool_call:raiseValueError(f"Expected at least one tool call, but got {len(tool_calls)} tool calls.")else:return[]tool_selections=[]fortool_callintool_calls:ifnotisinstance(tool_call,genai.protos.FunctionCall):raiseValueError("Invalid tool_call object")tool_selections.append(ToolSelection(tool_id=str(uuid.uuid4()),tool_name=tool_call.name,tool_kwargs=dict(tool_call.args),))returntool_selections@dispatcher.spandefstructured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->Model:"""Structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnsuper().structured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)@dispatcher.spanasyncdefastructured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->Model:"""Structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnawaitsuper().astructured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)@dispatcher.spandefstream_structured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->Generator[Union[Model,FlexibleModel],None,None]:"""Stream structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnsuper().stream_structured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)@dispatcher.spanasyncdefastream_structured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->AsyncGenerator[Union[Model,FlexibleModel],None]:"""Stream structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnawaitsuper().astream_structured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)
defget_tool_calls_from_response(self,response:ChatResponse,error_on_no_tool_call:bool=True,**kwargs:Any,)->List[ToolSelection]:"""Predict and call the tool."""tool_calls=response.message.additional_kwargs.get("tool_calls",[])iflen(tool_calls)<1:iferror_on_no_tool_call:raiseValueError(f"Expected at least one tool call, but got {len(tool_calls)} tool calls.")else:return[]tool_selections=[]fortool_callintool_calls:ifnotisinstance(tool_call,genai.protos.FunctionCall):raiseValueError("Invalid tool_call object")tool_selections.append(ToolSelection(tool_id=str(uuid.uuid4()),tool_name=tool_call.name,tool_kwargs=dict(tool_call.args),))returntool_selections
@dispatcher.spandefstructured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->Model:"""Structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnsuper().structured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)
@dispatcher.spanasyncdefastructured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->Model:"""Structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnawaitsuper().astructured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)
@dispatcher.spandefstream_structured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->Generator[Union[Model,FlexibleModel],None,None]:"""Stream structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnsuper().stream_structured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)
@dispatcher.spanasyncdefastream_structured_predict(self,output_cls:Type[Model],prompt:PromptTemplate,llm_kwargs:Optional[Dict[str,Any]]=None,**prompt_args:Any,)->AsyncGenerator[Union[Model,FlexibleModel],None]:"""Stream structured predict."""llm_kwargs=llm_kwargsor{}ifself._is_function_call_model:llm_kwargs["tool_choice"]=("required"if"tool_choice"notinllm_kwargselsellm_kwargs["tool_choice"])# by default structured prediction uses function calling to extract structured outputs# here we force tool_choice to be requiredreturnawaitsuper().astream_structured_predict(output_cls,prompt,llm_kwargs=llm_kwargs,**prompt_args)