create GlossaryChain for orchestrating chain functionality

src/chain.py

@@ -0,0 +1,54 @@
+import os
+import streamlit as st
+from langchain_community.vectorstores import FAISS
+from langchain_mistralai.embeddings import MistralAIEmbeddings
+from langchain_mistralai.chat_models import ChatMistralAI
+from langchain_core.prompts import PromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.documents import Document
+
+MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY")
+
+class GlossaryChain:
+    def __init__(self):
+        self.vector_store = load_vector_store()
+        self.retriever = self.vector_store.as_retriever()
+        self.llm = ChatMistralAI(
+            model="mistral-large-latest",
+            mistral_api_key=MISTRAL_API_KEY,
+            temperature=0.2
+        )
+        self.prompt = PromptTemplate.from_template(
+            "Answer the question based on the following context about Vanderbilt University Medical Center: \
+            \
+            {context}\n\nQuestion: \
+            \
+            {question}\n\nAnswer:"
+        )
+        self.chain = (
+            {"context": self.retriever | format_docs, "question": RunnablePassthrough()}
+            | self.prompt
+            | self.llm
+            | StrOutputParser()
+        )
+
+    def stream(self, input: str) -> str:
+        return self.chain.stream(input=input)
+
+    def invoke(self, input: str) -> str:
+        return self.chain.invoke(input=input)
+        
+
+def format_docs(docs: list[Document]) -> str:
+    """Format retrieved documents into a readable string"""
+    return "\n\n".join(doc.page_content for doc in docs)
+
+@st.cache_data
+def load_vector_store() -> FAISS:
+    return FAISS.load_local(
+        folder_path="faiss_index", 
+        embeddings = MistralAIEmbeddings(model="mistral-embed", mistral_api_key=MISTRAL_API_KEY),
+        allow_dangerous_deserialization=True
+    )
+