Create chain.py

chain.py

@@ -0,0 +1,53 @@
+from typing import List, Optional, Any
+from types import SimpleNamespace
+
+from langchain.prompts import PromptTemplate
+from langchain.chains.llm import LLMChain
+from langchain_core.outputs import ChatGeneration
+
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_openai import ChatOpenAI
+
+class GeminiChain(LLMChain):
+    def __init__(self, **kwargs) -> None:                
+        template = kwargs.pop('template')
+        model_name = kwargs.pop('model_name')
+        temperature = kwargs.pop('temperature', 0.3)
+        api_key = kwargs.pop('api_key')
+        verbose = kwargs.pop('verbose', False)
+        input_variables = kwargs.pop('input_variables', ['query',
+                                                         "history"])
+
+        output_parser = kwargs.pop('output_parser', None)
+        # create Gemini LLM object
+        # llm = ChatGoogleGenerativeAI(model=model_name,
+        #                                   temperature=temperature,
+        #                                   google_api_key=api_key)
+        llm = ChatOpenAI(api_key=api_key, model=model_name, temperature=temperature)
+        llm._get_llm_string()
+        # Define a prompt template
+        prompt_template = PromptTemplate(input_variables=input_variables,
+                                              template=template,
+                                              output_parser=output_parser)
+
+        super().__init__(llm=llm,
+                         prompt=prompt_template,
+                         verbose=verbose, **kwargs)
+
+    @classmethod
+    def from_system_prompt(self, template: str,
+                           input_variables: List[str],
+                           output_parser=None, **kwargs) -> LLMChain:
+        # create Gemini LLM object
+        llm_chain = self(template=template,
+                              input_variables=input_variables,
+                              output_parser=output_parser, **kwargs)
+        return llm_chain 
+    
+    def _get_llm_string(self, stop: Optional[List[str]] = None, **kwargs: Any) -> str:
+        return self._get_llm_string(stop=stop, **kwargs)
+    
+    def _infer(self, inputs: dict) -> ChatGeneration:
+        # Run LLM chain
+        llm_output = self.generate([inputs]).generations[0][0]
+        return llm_output