Update app.py

app.py

@@ -4,6 +4,9 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.chains import RetrievalQAWithSourcesChain
 from langchain.chat_models import ChatOpenAI
+from typing import Any, List, Mapping, Optional
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms.base import LLM
 from langchain.prompts.chat import (
     ChatPromptTemplate,
     SystemMessagePromptTemplate,
@@ -53,12 +56,55 @@ async def init():
     # docsearch = await cl.make_async(Chroma.from_documents)(pdf_data, embeddings) 
     docsearch = Chroma.from_documents(pdf_data, embeddings)
 
+    # custom SageMaker Model
+    class Llama2SageMaker(LLM):
+        max_new_tokens: int = 256
+        top_p: float = 0.9
+        temperature: float = 0.1
+    
+        @property
+        def _llm_type(self) -> str:
+            return "Llama2SageMaker"
+    
+        def _call(
+            self,
+            prompt: str,
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+        ) -> str:
+            if stop is not None:
+                raise ValueError("stop kwargs are not permitted.")
+            
+            json_body = {
+                "inputs" : [
+                  [{"role" : "user", "content" : prompt}]
+                ],
+                "parameters" : {
+                    "max_new_tokens" : self.max_new_tokens,
+                    "top_p" : self.top_p,
+                    "temperature" : self.temperature
+                }
+            }
+    
+            response = requests.post(model_api_gateway, json=json_body)
+    
+            return response.json()[0]["generation"]["content"]
+    
+        @property
+        def _identifying_params(self) -> Mapping[str, Any]:
+            """Get the identifying parameters."""
+            return {
+                "max_new_tokens" : self.max_new_tokens,
+                "top_p" : self.top_p,
+                "temperature" : self.temperature
+            }
+
+    # set our llm to the custom Llama2SageMaker endpoint model
+    llm = Llama2SageMaker()
+
     # Create a chain that uses the Chroma vector store
     chain = RetrievalQAWithSourcesChain.from_chain_type(
-        ChatOpenAI(
-            model_name="gpt-3.5-turbo-16k",
-            temperature=0,
-        ),
+        llm=llm,
         chain_type="stuff",
         retriever=docsearch.as_retriever(),
         return_source_documents=True,