"""Wrapper around IBM GENAI APIs for use in Langchain"""importasyncioimportloggingfromfunctoolsimportpartialfrompathlibimportPathfromtypingimportAny,Iterator,Optional,UnionfrompydanticimportConfigDictfrompydantic.v1importvalidatorfromgenaiimportClientfromgenai._utils.generalimportto_model_instance,to_model_optionalfromgenai.extensions._common.utilsimport(_prepare_generation_request,create_generation_info,create_generation_info_from_response,)fromgenai.schemaimport(ModerationParameters,PromptTemplateData,TextGenerationParameters,TextGenerationStreamCreateResponse,)fromgenai.text.generationimportCreateExecutionOptionstry:fromlangchain_core.callbacks.managerimport(AsyncCallbackManagerForLLMRun,CallbackManagerForLLMRun,)fromlangchain_core.language_models.llmsimportLLMfromlangchain_core.messagesimportBaseMessage,get_buffer_stringfromlangchain_core.outputsimportLLMResultfromgenai.extensions.langchain.utilsimport(CustomGenerationChunk,create_llm_output,dump_optional_model,load_config,update_token_usage,update_token_usage_stream,)exceptImportError:raiseImportError("Could not import langchain: Please install ibm-generative-ai[langchain] extension.")# noqa: B904logger=logging.getLogger(__name__)__all__=["LangChainInterface"]
[docs]classLangChainInterface(LLM):""" Class representing the LangChainChatInterface for interacting with the LangChain chat API. Example:: from genai import Client, Credentials from genai.extensions.langchain import LangChainInterface from genai.schema import TextGenerationParameters client = Client(credentials=Credentials.from_env()) llm = LangChainInterface( client=client, model_id="meta-llama/llama-3-1-70b-instruct", parameters=TextGenerationParameters(max_new_tokens=50) ) response = chat_model.generate(prompts=["Hello world!"]) print(response) """model_config=ConfigDict(extra="forbid",protected_namespaces=())client:Clientmodel_id:strprompt_id:Optional[str]=Noneparameters:Optional[TextGenerationParameters]=Nonemoderations:Optional[ModerationParameters]=Nonedata:Optional[PromptTemplateData]=Nonestreaming:Optional[bool]=Noneexecution_options:Optional[CreateExecutionOptions]=None@validator("parameters","moderations","data","execution_options",pre=True,always=True)@classmethoddef_validate_data_models(cls,value,values,config,field):returnto_model_optional(value,Model=field.type_,copy=False)@propertydef_common_identifying_params(self):return{"model_id":self.model_id,"prompt_id":self.prompt_id,"parameters":dump_optional_model(self.parameters),"moderations":dump_optional_model(self.moderations),"data":dump_optional_model(self.data),**super()._identifying_params,}@propertydef_identifying_stream_params(self)->dict[str,Any]:returnself._common_identifying_params@propertydef_identifying_params(self)->dict[str,Any]:return{**self._common_identifying_params,"execution_options":dump_optional_model(self.execution_options),}
@propertydef_llm_type(self)->str:"""Return type of llm."""return"ibmgenai_llm"def_prepare_request(self,**kwargs):updated={k:kwargs.pop(k,v)fork,vinself._identifying_params.items()}return_prepare_generation_request(**kwargs,**updated)def_prepare_stream_request(self,**kwargs):updated={k:kwargs.pop(k,v)fork,vinself._identifying_stream_params.items()}return_prepare_generation_request(**kwargs,**updated)def_call(self,prompt:str,stop:Optional[list[str]]=None,run_manager:Optional[CallbackManagerForLLMRun]=None,**kwargs:Any,)->str:result=self._generate(prompts=[prompt],stop=stop,run_manager=run_manager,**kwargs)returnresult.generations[0][0].textdef_generate(self,prompts:list[str],stop:Optional[list[str]]=None,run_manager:Optional[CallbackManagerForLLMRun]=None,**kwargs:Any,)->LLMResult:final_result=LLMResult(generations=[],llm_output=create_llm_output(model=self.model_id))assertfinal_result.llm_outputiflen(prompts)==0:returnfinal_resultifself.streaming:iflen(prompts)!=1:raiseValueError("Streaming works only for a single prompt.")generation=CustomGenerationChunk(text="",generation_info=create_generation_info())forchunk_resultinself._stream(prompt=prompts[0],stop=stop,run_manager=run_manager,**kwargs,):assertchunk_result.generation_infotoken_usage=chunk_result.generation_info.pop("token_usage")generation+=chunk_resultassertgeneration.generation_infoupdate_token_usage_stream(target=generation.generation_info["token_usage"],source=token_usage,)final_result.generations.append([generation])assertgeneration.generation_infoupdate_token_usage(target=final_result.llm_output["token_usage"],source=generation.generation_info["token_usage"])returnfinal_resultelse:responses=list(self.client.text.generation.create(**self._prepare_request(inputs=prompts,stop=stop,**kwargs)))forresponseinresponses:forresultinresponse.results:generation_info=create_generation_info_from_response(response,result=result)chunk=CustomGenerationChunk(text=result.generated_textor"",generation_info=generation_info,)logger.info("Output of GENAI call: {}".format(chunk.text))update_token_usage(target=final_result.llm_output["token_usage"],source=generation_info["token_usage"])final_result.generations.append([chunk])returnfinal_resultasyncdef_agenerate(self,prompts:list[str],stop:Optional[list[str]]=None,run_manager:Optional[AsyncCallbackManagerForLLMRun]=None,**kwargs:Any,)->LLMResult:returnawaitasyncio.get_running_loop().run_in_executor(None,partial(self._generate,prompts,stop,run_manager,**kwargs))def_stream(self,prompt:str,stop:Optional[list[str]]=None,run_manager:Optional[CallbackManagerForLLMRun]=None,**kwargs:Any,)->Iterator[CustomGenerationChunk]:params=to_model_instance(self.parameters,TextGenerationParameters)params.stop_sequences=stoporparams.stop_sequencesdefsend_chunk(*,text:Optional[str]=None,generation_info:dict,response:TextGenerationStreamCreateResponse):logger.info("Chunk received: {}".format(text))chunk=CustomGenerationChunk(text=textor"",generation_info=generation_info.copy(),)yieldchunkifrun_manager:run_manager.on_llm_new_token(token=chunk.text,chunk=chunk,response=response)forresponseinself.client.text.generation.create_stream(**self._prepare_stream_request(input=prompt,stop=stop,**kwargs)):ifresponse.moderations:generation_info=create_generation_info_from_response(response,result=response.moderation)yield fromsend_chunk(generation_info=generation_info,response=response)forresultinresponse.resultsor[]:generation_info=create_generation_info_from_response(response,result=result)yield fromsend_chunk(text=result.generated_text,generation_info=generation_info,response=response)