first commit

.chainlit/config.toml

@@ -0,0 +1,84 @@
+[project]
+# Whether to enable telemetry (default: true). No personal data is collected.
+enable_telemetry = true
+
+# List of environment variables to be provided by each user to use the app.
+user_env = []
+
+# Duration (in seconds) during which the session is saved when the connection is lost
+session_timeout = 3600
+
+# Enable third parties caching (e.g LangChain cache)
+cache = false
+
+# Follow symlink for asset mount (see https://github.com/Chainlit/chainlit/issues/317)
+# follow_symlink = false
+
+[features]
+# Show the prompt playground
+prompt_playground = false
+
+# Process and display HTML in messages. This can be a security risk (see https://stackoverflow.com/questions/19603097/why-is-it-dangerous-to-render-user-generated-html-or-javascript)
+unsafe_allow_html = false
+
+# Process and display mathematical expressions. This can clash with "$" characters in messages.
+latex = false
+
+# Authorize users to upload files with messages
+multi_modal = false
+
+# Allows user to use speech to text
+[features.speech_to_text]
+    enabled = false
+    # See all languages here https://github.com/JamesBrill/react-speech-recognition/blob/HEAD/docs/API.md#language-string
+    # language = "en-US"
+
+[UI]
+# Name of the app and chatbot.
+name = "Sage"
+
+# Show the readme while the conversation is empty.
+show_readme_as_default = false
+
+# Description of the app and chatbot. This is used for HTML tags.
+# description = ""
+
+# Large size content are by default collapsed for a cleaner ui
+default_collapse_content = true
+
+# The default value for the expand messages settings.
+default_expand_messages = false
+
+# Hide the chain of thought details from the user in the UI.
+hide_cot = false
+
+# Link to your github repo. This will add a github button in the UI's header.
+# github = ""
+
+# Specify a CSS file that can be used to customize the user interface.
+# The CSS file can be served from the public directory or via an external link.
+# custom_css = "/public/test.css"
+
+# Override default MUI light theme. (Check theme.ts)
+[UI.theme.light]
+    #background = "#FAFAFA"
+    #paper = "#FFFFFF"
+
+    [UI.theme.light.primary]
+        #main = "#F80061"
+        #dark = "#980039"
+        #light = "#FFE7EB"
+
+# Override default MUI dark theme. (Check theme.ts)
+[UI.theme.dark]
+    #background = "#FAFAFA"
+    #paper = "#FFFFFF"
+
+    [UI.theme.dark.primary]
+        #main = "#F80061"
+        #dark = "#980039"
+        #light = "#FFE7EB"
+
+
+[meta]
+generated_by = "0.7.700"

.chainlit/translations/en-US.json

@@ -0,0 +1,231 @@
+{
+    "components": {
+        "atoms": {
+            "buttons": {
+                "userButton": {
+                    "menu": {
+                        "settings": "Settings",
+                        "settingsKey": "S",
+                        "APIKeys": "API Keys",
+                        "logout": "Logout"
+                    }
+                }
+            }
+        },
+        "molecules": {
+            "newChatButton": {
+                "newChat": "New Chat"
+            },
+            "tasklist": {
+                "TaskList": {
+                    "title": "\ud83d\uddd2\ufe0f Task List",
+                    "loading": "Loading...",
+                    "error": "An error occured"
+                }
+            },
+            "attachments": {
+                "cancelUpload": "Cancel upload",
+                "removeAttachment": "Remove attachment"
+            },
+            "newChatDialog": {
+                "createNewChat": "Create new chat?",
+                "clearChat": "This will clear the current messages and start a new chat.",
+                "cancel": "Cancel",
+                "confirm": "Confirm"
+            },
+            "settingsModal": {
+                "settings": "Settings",
+                "expandMessages": "Expand Messages",
+                "hideChainOfThought": "Hide Chain of Thought",
+                "darkMode": "Dark Mode"
+            },
+            "detailsButton": {
+                "using": "Using",
+                "running": "Running",
+                "took_one": "Took {{count}} step",
+                "took_other": "Took {{count}} steps"
+            },
+            "auth": {
+                "authLogin": {
+                    "title": "Login to access the app.",
+                    "form": {
+                        "email": "Email address",
+                        "password": "Password",
+                        "noAccount": "Don't have an account?",
+                        "alreadyHaveAccount": "Already have an account?",
+                        "signup": "Sign Up",
+                        "signin": "Sign In",
+                        "or": "OR",
+                        "continue": "Continue",
+                        "forgotPassword": "Forgot password?",
+                        "passwordMustContain": "Your password must contain:",
+                        "emailRequired": "email is a required field",
+                        "passwordRequired": "password is a required field"
+                    },
+                    "error": {
+                        "default": "Unable to sign in.",
+                        "signin": "Try signing in with a different account.",
+                        "oauthsignin": "Try signing in with a different account.",
+                        "redirect_uri_mismatch": "The redirect URI is not matching the oauth app configuration.",
+                        "oauthcallbackerror": "Try signing in with a different account.",
+                        "oauthcreateaccount": "Try signing in with a different account.",
+                        "emailcreateaccount": "Try signing in with a different account.",
+                        "callback": "Try signing in with a different account.",
+                        "oauthaccountnotlinked": "To confirm your identity, sign in with the same account you used originally.",
+                        "emailsignin": "The e-mail could not be sent.",
+                        "emailverify": "Please verify your email, a new email has been sent.",
+                        "credentialssignin": "Sign in failed. Check the details you provided are correct.",
+                        "sessionrequired": "Please sign in to access this page."
+                    }
+                },
+                "authVerifyEmail": {
+                    "almostThere": "You're almost there! We've sent an email to ",
+                    "verifyEmailLink": "Please click on the link in that email to complete your signup.",
+                    "didNotReceive": "Can't find the email?",
+                    "resendEmail": "Resend email",
+                    "goBack": "Go Back",
+                    "emailSent": "Email sent successfully.",
+                    "verifyEmail": "Verify your email address"
+                },
+                "providerButton": {
+                    "continue": "Continue with {{provider}}",
+                    "signup": "Sign up with {{provider}}"
+                },
+                "authResetPassword": {
+                    "newPasswordRequired": "New password is a required field",
+                    "passwordsMustMatch": "Passwords must match",
+                    "confirmPasswordRequired": "Confirm password is a required field",
+                    "newPassword": "New password",
+                    "confirmPassword": "Confirm password",
+                    "resetPassword": "Reset Password"
+                },
+                "authForgotPassword": {
+                    "email": "Email address",
+                    "emailRequired": "email is a required field",
+                    "emailSent": "Please check the email address {{email}} for instructions to reset your password.",
+                    "enterEmail": "Enter your email address and we will send you instructions to reset your password.",
+                    "resendEmail": "Resend email",
+                    "continue": "Continue",
+                    "goBack": "Go Back"
+                }
+            }
+        },
+        "organisms": {
+            "chat": {
+                "history": {
+                    "index": {
+                        "showHistory": "Show history",
+                        "lastInputs": "Last Inputs",
+                        "noInputs": "Such empty...",
+                        "loading": "Loading..."
+                    }
+                },
+                "inputBox": {
+                    "input": {
+                        "placeholder": "Type your message here..."
+                    },
+                    "speechButton": {
+                        "start": "Start recording",
+                        "stop": "Stop recording"
+                    },
+                    "SubmitButton": {
+                        "sendMessage": "Send message",
+                        "stopTask": "Stop Task"
+                    },
+                    "UploadButton": {
+                        "attachFiles": "Attach files"
+                    },
+                    "waterMark": {
+                        "text": "Built with"
+                    }
+                },
+                "Messages": {
+                    "index": {
+                        "running": "Running",
+                        "executedSuccessfully": "executed successfully",
+                        "failed": "failed",
+                        "feedbackUpdated": "Feedback updated",
+                        "updating": "Updating"
+                    }
+                },
+                "dropScreen": {
+                    "dropYourFilesHere": "Drop your files here"
+                },
+                "index": {
+                    "failedToUpload": "Failed to upload",
+                    "cancelledUploadOf": "Cancelled upload of",
+                    "couldNotReachServer": "Could not reach the server",
+                    "continuingChat": "Continuing previous chat"
+                },
+                "settings": {
+                    "settingsPanel": "Settings panel",
+                    "reset": "Reset",
+                    "cancel": "Cancel",
+                    "confirm": "Confirm"
+                }
+            },
+            "threadHistory": {
+                "sidebar": {
+                    "filters": {
+                        "FeedbackSelect": {
+                            "feedbackAll": "Feedback: All",
+                            "feedbackPositive": "Feedback: Positive",
+                            "feedbackNegative": "Feedback: Negative"
+                        },
+                        "SearchBar": {
+                            "search": "Search"
+                        }
+                    },
+                    "DeleteThreadButton": {
+                        "confirmMessage": "This will delete the thread as well as it's messages and elements.",
+                        "cancel": "Cancel",
+                        "confirm": "Confirm",
+                        "deletingChat": "Deleting chat",
+                        "chatDeleted": "Chat deleted"
+                    },
+                    "index": {
+                        "pastChats": "Past Chats"
+                    },
+                    "ThreadList": {
+                        "empty": "Empty...",
+                        "today": "Today",
+                        "yesterday": "Yesterday",
+                        "previous7days": "Previous 7 days",
+                        "previous30days": "Previous 30 days"
+                    },
+                    "TriggerButton": {
+                        "closeSidebar": "Close sidebar",
+                        "openSidebar": "Open sidebar"
+                    }
+                },
+                "Thread": {
+                    "backToChat": "Go back to chat",
+                    "chatCreatedOn": "This chat was created on"
+                }
+            },
+            "header": {
+                "chat": "Chat",
+                "readme": "Readme"
+            }
+        }
+    },
+    "hooks": {
+        "useLLMProviders": {
+            "failedToFetchProviders": "Failed to fetch providers:"
+        }
+    },
+    "pages": {
+        "Design": {},
+        "Env": {
+            "savedSuccessfully": "Saved successfully",
+            "requiredApiKeys": "Required API Keys",
+            "requiredApiKeysInfo": "To use this app, the following API keys are required. The keys are stored on your device's local storage."
+        },
+        "Page": {
+            "notPartOfProject": "You are not part of this project."
+        },
+        "ResumeButton": {
+            "resumeChat": "Resume Chat"
+        }
+    }
+}

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.github/workflows/check_file_size.yml

@@ -0,0 +1,16 @@
+name: Check file size
+on:               # or directly `on: [push]` to run the action on every push on any branch
+  pull_request:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check large files
+        uses: ActionsDesk/lfs-warning@v2.0
+        with:
+          filesizelimit: 10485760 # this is 10MB so we can sync to HF Spaces

.github/workflows/sync_to_hf_space.yml

@@ -0,0 +1,18 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push -f https://nav13n:$HF_TOKEN@huggingface.co/spaces/nav13n/DocuSage main

.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.11
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+WORKDIR $HOME/app
+COPY --chown=user . $HOME/app
+COPY ./requirements.txt ~/app/requirements.txt
+RUN pip install -r requirements.txt
+COPY . .
+CMD ["chainlit", "run", "app.py", "--port", "7860"]

README.md

@@ -0,0 +1,14 @@
+---
+title: DocuSage
+emoji: ⚡
+colorFrom: indigo
+colorTo: gray
+sdk: docker
+pinned: false
+license: mit
+short_description: DocuSage -AI-powered insights in seconds. 
+---
+
+# DocuSage
+
+DocuSage is a powerful document question-answering application that allows users to upload documents and ask natural language questions about their content. With cutting-edge AI technology, DocuSage provides accurate, concise answers in seconds, saving time and boosting productivity. Unlock valuable insights and let DocuSage be your sage guide through the world of document knowledge.

app.py

@@ -0,0 +1,30 @@
+import os
+import chainlit as cl 
+from sage import app
+from langchain.schema.runnable import RunnableConfig
+
+
+welcome_message = "Welcome! Ask anything about your stored documents and get AI-powered insights in seconds."
+
+@cl.on_chat_start  
+async def start_chat():
+    await cl.Message(content=welcome_message).send()
+    cl.user_session.set("runnable", app)
+
+
+@cl.on_message 
+async def main(message: cl.Message):
+    runnable = cl.user_session.get("runnable")
+    msg = cl.Message(content="")
+
+    input = {"question": message.content}
+
+    value = None
+    for output in runnable.stream(input):
+        for key, value in output.items():
+            print(f"Finished running: {key}:")
+            if key == "generator_agent":
+                answer = value["answer"]
+                await msg.stream_token(answer)
+
+    await msg.send()

chain.py

@@ -0,0 +1,71 @@
+# Langchain imports
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Qdrant
+from langchain.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+from langchain_openai import OpenAIEmbeddings
+from langchain_core.output_parsers import StrOutputParser
+from langchain.schema.runnable import RunnablePassthrough
+from dotenv import load_dotenv
+
+
+load_dotenv()
+
+######################## Build RAG Chain #############################
+######################################################################
+
+#### Load  Documents
+loader = PyMuPDFLoader(
+    "https://d18rn0p25nwr6d.cloudfront.net/CIK-0001326801/c7318154-f6ae-4866-89fa-f0c589f2ee3d.pdf"
+)
+
+documents = loader.load()
+
+#### Split Documents
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size = 800,
+    chunk_overlap = 100
+)
+
+documents = text_splitter.split_documents(documents)
+
+embeddings = OpenAIEmbeddings(
+    model="text-embedding-3-small"
+)
+
+###  Create Vector Store
+vector_store = Qdrant.from_documents(
+    documents,
+    embeddings,
+    location=":memory:",
+    collection_name="Meta 10k Filings",
+)
+
+### Create Prmopt Template
+template = """Answer the question based only on the following context. If you cannot answer the question with the context, please respond with 'I don't know':
+
+Context:
+{context}
+
+Question:
+{question}
+"""
+prompt = ChatPromptTemplate.from_template(template)
+
+def format_docs(docs):
+    return "\n\n".join([d.page_content for d in docs])
+
+### Setup RAG Chain
+
+retriever = vector_store.as_retriever(search_type="similarity_score_threshold", 
+                                      search_kwargs={"score_threshold": 0.6, "k":8})
+primary_qa_llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
+
+rag_chain = (
+    {"context": retriever | format_docs, "question": RunnablePassthrough()}
+    | prompt 
+    | primary_qa_llm
+    | StrOutputParser()
+
+)

chainlit.md

@@ -0,0 +1,3 @@
+# DocuSage
+
+DocuSage is a powerful document question-answering application that allows users to upload documents and ask natural language questions about their content. With cutting-edge AI technology, DocuSage provides accurate, concise answers in seconds, saving time and boosting productivity. Unlock valuable insights and let DocuSage be your sage guide through the world of document knowledge.

environment.yaml

@@ -0,0 +1,8 @@
+name: sage
+
+channels:
+  - conda-forge
+  - defaults
+
+dependencies:
+  - python=3.11

requirements.txt

@@ -0,0 +1,15 @@
+chainlit
+tiktoken
+python-dotenv
+langchain
+langchain_core
+langchain-community
+langchainhub
+faiss-cpu
+langgraph
+fastembed
+flashrank
+pandas
+huggingface-hub
+langchain-groq
+pymupdf

sage.py

@@ -0,0 +1,684 @@
+
+from langchain_community.vectorstores import FAISS
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings.fastembed import FastEmbedEmbeddings
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
+embed_model = FastEmbedEmbeddings(model_name="snowflake/snowflake-arctic-embed-m")
+
+from groq import Groq
+from langchain_groq import ChatGroq
+
+
+llm = ChatGroq(temperature=0,
+                      model_name="Llama3-8b-8192",
+                      api_key=os.getenv("GROQ_API_KEY"),)
+
+loader = PyMuPDFLoader("https://home.synise.com/HRUtility/Documents/HRA/UmaP/Synise%20Handbook.pdf")
+documents = loader.load()
+
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=1500, chunk_overlap=200
+)
+doc_splits = text_splitter.split_documents(documents)
+
+print(len(doc_splits),doc_splits[0])
+
+vectorstore = FAISS.from_documents(documents=doc_splits,embedding=embed_model)
+
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import FlashrankRerank
+
+compressor = FlashrankRerank()
+retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 20})
+compression_retriever = ContextualCompressionRetriever(
+    base_compressor=compressor, base_retriever=retriever
+)
+
+from operator import itemgetter
+from langchain.prompts import PromptTemplate
+from langchain.schema.runnable import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+
+
+RAG_PROMPT_TEMPLATE = """
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+ Answer the question based only on the provided context. If you cannot answer the question with the provided context, please respond with 'I don't know" without any preamble, explanation, or additional text.
+
+Context:
+{context}
+
+Question:
+{question}
+
+<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+
+rag_prompt = PromptTemplate(
+    template=RAG_PROMPT_TEMPLATE, input_variables=["question", "context"]
+)
+
+response_chain = (rag_prompt
+    | llm
+    | StrOutputParser()
+
+)
+
+def dummy_payroll_api_call(employee_id, month, year):
+
+  data = {
+    2023: {
+        "MAY": {
+            "employeeDetails": {
+                "employeeId": "E2468",
+                "firstName": "Sarah",
+                "lastName": "Thompson",
+                "designation": "Product Manager"
+            },
+            "paymentDetails": {
+                "year": 2023,
+                "month": "JAN",
+                "basicSalary": 5500,
+                "allowances": [
+                    {
+                        "type": "Housing Allowance",
+                        "amount": 1500
+                    },
+                    {
+                        "type": "Travel Allowance",
+                        "amount": 800
+                    }
+                ],
+                "deductions": [
+                    {
+                        "type": "Provident Fund",
+                        "amount": 650
+                    },
+                    {
+                        "type": "Health Insurance",
+                        "amount": 300
+                    }
+                ],
+                "taxes": [
+                    {
+                        "type": "Income Tax",
+                        "amount": 1300
+                    }
+                ],
+                "grossSalary": 7800,
+                "totalDeductions": 2250,
+                "netSalary": 6650
+            },
+            "companyDetails": {
+                "companyName": "Tech Solutions Ltd.",
+                "address": "789 Maple Avenue, City"
+            }
+        }
+    },
+    2024: {
+        "JAN": {
+            "employeeDetails": {
+                "employeeId": "E2468",
+                "firstName": "Sarah",
+                "lastName": "Thompson",
+                "designation": "Product Manager"
+            },
+            "paymentDetails": {
+                "year": 2024,
+                "month": "JAN",
+                "basicSalary": 6500,
+                "allowances": [
+                    {
+                        "type": "Housing Allowance",
+                        "amount": 1500
+                    },
+                    {
+                        "type": "Travel Allowance",
+                        "amount": 800
+                    }
+                ],
+                "deductions": [
+                    {
+                        "type": "Provident Fund",
+                        "amount": 650
+                    },
+                    {
+                        "type": "Health Insurance",
+                        "amount": 300
+                    }
+                ],
+                "taxes": [
+                    {
+                        "type": "Income Tax",
+                        "amount": 1300
+                    }
+                ],
+                "grossSalary": 8800,
+                "totalDeductions": 2250,
+                "netSalary": 6550
+            },
+            "companyDetails": {
+                "companyName": "Tech Solutions Ltd.",
+                "address": "789 Maple Avenue, City"
+            }
+        },
+        "FEB": {
+            "employeeDetails": {
+                "employeeId": "E2468",
+                "firstName": "Sarah",
+                "lastName": "Thompson",
+                "designation": "Product Manager"
+            },
+            "paymentDetails": {
+                "year": 2024,
+                "month": "FEB",
+                "basicSalary": 6500,
+                "allowances": [
+                    {
+                        "type": "Housing Allowance",
+                        "amount": 1500
+                    },
+                    {
+                        "type": "Travel Allowance",
+                        "amount": 800
+                    }
+                ],
+                "deductions": [
+                    {
+                        "type": "Provident Fund",
+                        "amount": 650
+                    },
+                    {
+                        "type": "Health Insurance",
+                        "amount": 300
+                    }
+                ],
+                "taxes": [
+                    {
+                        "type": "Income Tax",
+                        "amount": 1300
+                    }
+                ],
+                "grossSalary": 8800,
+                "totalDeductions": 2250,
+                "netSalary": 6550
+            },
+            "companyDetails": {
+                "companyName": "Tech Solutions Ltd.",
+                "address": "789 Maple Avenue, City"
+            }
+        },
+                "MAY": {
+            "employeeDetails": {
+                "employeeId": "E2468",
+                "firstName": "Sarah",
+                "lastName": "Thompson",
+                "designation": "Product Manager"
+            },
+            "paymentDetails": {
+                "year": 2024,
+                "month": "MAY",
+                "basicSalary": 6500,
+                "allowances": [
+                    {
+                        "type": "Housing Allowance",
+                        "amount": 1500
+                    },
+                    {
+                        "type": "Travel Allowance",
+                        "amount": 800
+                    }
+                ],
+                "deductions": [
+                    {
+                        "type": "Provident Fund",
+                        "amount": 650
+                    },
+                    {
+                        "type": "Health Insurance",
+                        "amount": 300
+                    }
+                ],
+                "taxes": [
+                    {
+                        "type": "Income Tax",
+                        "amount": 1500
+                    }
+                ],
+                "grossSalary": 8800,
+                "totalDeductions": 2450,
+                "netSalary": 6350
+            },
+            "companyDetails": {
+                "companyName": "Tech Solutions Ltd.",
+                "address": "789 Maple Avenue, City"
+            }
+        },
+        "APR": {
+            "employeeDetails": {
+                "employeeId": "E2468",
+                "firstName": "Sarah",
+                "lastName": "Thompson",
+                "designation": "Product Manager"
+            },
+            "paymentDetails": {
+                "year": 2024,
+                "month": "APR",
+                "basicSalary": 6500,
+                "allowances": [
+                    {
+                        "type": "Housing Allowance",
+                        "amount": 1500
+                    },
+                    {
+                        "type": "Travel Allowance",
+                        "amount": 800
+                    }
+                ],
+                "deductions": [
+                    {
+                        "type": "Provident Fund",
+                        "amount": 650
+                    },
+                    {
+                        "type": "Health Insurance",
+                        "amount": 300
+                    }
+                ],
+                "taxes": [
+                    {
+                        "type": "Income Tax",
+                        "amount": 1500
+                    }
+                ],
+                "grossSalary": 8800,
+                "totalDeductions": 2450,
+                "netSalary": 6350
+            },
+            "companyDetails": {
+                "companyName": "Tech Solutions Ltd.",
+                "address": "789 Maple Avenue, City"
+            }
+        }
+    }
+}
+  year= 2024 if year == "CUR" else year
+  year= 2023 if year == "PREV" else year
+
+  month= "MAY" if month == "CUR" else month
+  month= "APR" if month == "PREV" else month
+
+
+  return data[year][month]
+
+print(dummy_payroll_api_call(1234, 'CUR', 2024))
+
+import time
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from langchain_core.output_parsers import StrOutputParser
+
+ROUTER_AGENT_PROMPT_TEMPLATE = """
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+You are an expert at delegating user questions to one of the most appropriate agents 'policy_agent' or 'payroll_agent'.
+
+Use the following criteria to determine the appropriate agents to answer the user que:
+
+- If the query is regarding payslips, salary, tax deductions, basepay of a given month, use payroll_agent'.
+- If the question is closely related to general human resource queries, organisational policies, prompt engineering, or adversarial attacks, even if the keywords are not explicitly mentioned, use the 'policyagent'.
+
+Your output should be a JSON object with a single key 'agent' and a value of either 'policy_agent' or 'payroll_agent'. Do not include any preamble, explanation, or additional text.
+
+User's Question: {question}
+
+<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+
+router_prompt = PromptTemplate(
+    template=ROUTER_AGENT_PROMPT_TEMPLATE, input_variables=["question"]
+)
+
+
+router_chain = router_prompt | llm | JsonOutputParser()
+
+print(router_chain.invoke({"question":"What is my salary on 6 2024 ?"}))
+
+print(router_chain.invoke({"question":"What is leave policy ?"}))
+
+payroll_schema= {
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "Monthly Payslip",
+  "description": "A schema for a monthly payslip",
+  "type": "object",
+  "properties": {
+    "employeeDetails": {
+      "type": "object",
+      "properties": {
+        "employeeId": {
+          "type": "string",
+          "description": "Unique identifier for the employee"
+        },
+        "firstName": {
+          "type": "string",
+          "description": "First name of the employee"
+        },
+        "lastName": {
+          "type": "string",
+          "description": "Last name of the employee"
+        },
+        "designation": {
+          "type": "string",
+          "description": "Designation or job title of the employee"
+        }
+      },
+      "required": ["employeeId", "firstName", "lastName", "designation"]
+    },
+    "paymentDetails": {
+      "type": "object",
+      "properties": {
+        "year": {
+          "type": "integer",
+          "description": "Year of the pay period"
+        },
+        "month": {
+          "type": "string",
+          "enum": ["JAN", "FEB", "MAR", "APR", "MAY", "JUN", "JUL", "AUG", "SEP", "OCT", "NOV", "DEC"],
+          "description": "Month of the pay period"
+        },
+        "basicSalary": {
+          "type": "number",
+          "description": "Basic salary of the employee"
+        },
+        "allowances": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "properties": {
+              "type": {
+                "type": "string",
+                "description": "Type of allowance"
+              },
+              "amount": {
+                "type": "number",
+                "description": "Amount of the allowance"
+              }
+            },
+            "required": ["type", "amount"]
+          }
+        },
+        "deductions": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "properties": {
+              "type": {
+                "type": "string",
+                "description": "Type of deduction"
+              },
+              "amount": {
+                "type": "number",
+                "description": "Amount of the deduction"
+              }
+            },
+            "required": ["type", "amount"]
+          }
+        },
+        "taxes": {
+          "type": "array",
+          "items": {
+            "type": "object",
+            "properties": {
+              "type": {
+                "type": "string",
+                "description": "Type of tax"
+              },
+              "amount": {
+                "type": "number",
+                "description": "Amount of the tax"
+              }
+            },
+            "required": ["type", "amount"]
+          }
+        },
+        "grossSalary": {
+          "type": "number",
+          "description": "Gross salary (basic salary + allowances)"
+        },
+        "totalDeductions": {
+          "type": "number",
+          "description": "Total deductions (including taxes)"
+        },
+        "netSalary": {
+          "type": "number",
+          "description": "Net salary (gross salary - total deductions)"
+        }
+      },
+      "required": ["year", "month", "basicSalary", "allowances", "deductions", "taxes", "grossSalary", "totalDeductions", "netSalary"]
+    },
+    "companyDetails": {
+      "type": "object",
+      "properties": {
+        "companyName": {
+          "type": "string",
+          "description": "Name of the company"
+        },
+        "address": {
+          "type": "string",
+          "description": "Address of the company"
+        }
+      },
+      "required": ["companyName", "address"]
+    }
+  },
+  "required": ["employeeDetails", "paymentDetails", "companyDetails"]
+}
+
+print(str(payroll_schema))
+
+import time
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from langchain_core.output_parsers import StrOutputParser
+
+FILTER_EXTTRACTION_PROMPT = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+Extract the month and year from a given user question about payroll. Use the following schema instructions to guide your extraction.
+
+Instructions:
+1. Your output should be a JSON object with only two keys, 'month' and 'year'.
+2. 'month' key shall have value ["JAN", "FEB", "MAR", "APR", "MAY", "JUN", "JUL", "AUG", "SEP", "OCT", "NOV", "DEC"]
+3. 'year' shall be a number between 2020 and 2024.
+4. If the user is suggesting current year or month, respond with "CUR" for 'month' and 'year' keys accordingly
+5. If the user is suggesting previous year or month, respond with "PREV" for 'month' and 'year' keys accordingly
+
+
+Do not include any preamble, explanation, or additional text.
+
+User Question: {question}
+<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+
+filter_extraction_prompt = PromptTemplate(
+    template=FILTER_EXTTRACTION_PROMPT, input_variables=["question"]
+)
+
+fiter_extraction_chain = filter_extraction_prompt | llm | JsonOutputParser()
+
+print(fiter_extraction_chain.invoke({"question":"What is my salary on 6 2024 ?"}))
+
+import time
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from langchain_core.output_parsers import StrOutputParser
+
+PAYROLL_QA_PROMPT = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Answer the user query given the provided payroll data in json form. Use the  provided schema to understand the payroll data structure. If you cannot answer the question with the provided information, please respond with 'I don't know" without any preamble, explanation, or additional text
+
+SCHEMA:
+{schema}
+
+PAYROLL DATA
+{data}
+
+PAYROLL DATA:
+{data}
+
+User Question: {question}
+<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+
+payroll_qa_prompt = PromptTemplate(
+    template=PAYROLL_QA_PROMPT, input_variables=["question", "data", "schema"]
+)
+
+payroll_qa_chain = payroll_qa_prompt | llm | StrOutputParser()
+
+result = fiter_extraction_chain.invoke({"question":"What is my salary on jan 2024 ?"})
+
+result
+
+api_result = dummy_payroll_api_call(1234, result["month"], result["year"])
+
+api_result
+
+payroll_qa_chain.invoke({"question":"What is my salary on jan 2024 ?", "data":api_result, "schema":payroll_schema})
+
+from typing_extensions import TypedDict
+from typing import List
+
+### State
+
+class AgentState(TypedDict):
+    question : str
+    answer : str
+    documents : List[str]
+
+import logging as log
+
+def route_question(state):
+    """
+    Route question to payroll_agent or policy_agent to retrieve reevant data
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        str: Next node to call
+    """
+
+    question = state["question"]
+    result = router_chain.invoke({"question": question})
+
+    log.debug('Routing to {}....'.format(result["agent"]))
+
+    if result['agent'] == 'payroll_agent':
+        log.debug('Routing to {}....'.format(result["agent"]))
+        return "payroll_agent"
+    elif result['agent'] == 'policy_agent':
+        log.debug('Routing to {}....'.format(result["agent"]))
+        return "policy_agent"
+
+state = AgentState(question="What is my salary on jan 2024 ?", answer="", documents=None)
+route_question(state)
+
+from langchain.schema import Document
+def retrieve_policy(state):
+    """
+    Retrieve policy documents from vectorstore
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): New key added to state, documents, that contains retrieved documents
+    """
+    log.debug("Retreiving policy documents.......")
+    question = state["question"]
+    documents = compression_retriever.invoke(question)
+    return {"documents": documents, "question": question}
+
+state = AgentState(question="What is leave policy?", answer="", documents=None)
+retrieve_policy(state)
+
+def generate_answer(state):
+    """
+    Generate answer using retrieved data
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): New key added to state, generation, that contains LLM generation
+    """
+    log.debug("Generating answer.......")
+    question = state["question"]
+    documents = state["documents"]
+
+
+    answer = response_chain.invoke({"context": documents, "question": question})
+
+    return {"documents": documents, "question": question, "answer": answer}
+
+state = AgentState(question="What is leave policy?", answer="", documents=[Document(page_content="According to leave policy, there are two types of leaves 1: PL 2: CL")])
+generate_answer(state)
+
+def query_payroll(state):
+    """
+    Query payroll api to retrieve payroll data
+
+    Args:
+        state (dict): The current graph state
+
+    Returns:
+        state (dict): Updated state with retrived payroll data
+    """
+
+
+    question = state["question"]
+    payroll_query_filters = fiter_extraction_chain.invoke({"question":question})
+    payroll_api_query_results = dummy_payroll_api_call(1234, result["month"], result["year"])
+
+
+    context = context = 'PAYROLL DATA SCHEMA: \n {payroll_schema} \n PAYROLL DATA: {payroll_api_query_results}'.format(
+    payroll_schema=payroll_schema, payroll_api_query_results=payroll_api_query_results)
+
+    documents = [Document(page_content=context)]
+    return {"documents": documents, "question": question}
+
+state = AgentState(question="Tell me salary for Jan 2024?", answer="", documents=None)
+query_payroll(state)
+
+from langgraph.graph import END, StateGraph
+workflow = StateGraph(AgentState)
+
+# Define the nodes
+workflow.add_node("payroll_agent", query_payroll)
+workflow.add_node("policy_agent", retrieve_policy)
+workflow.add_node("generator_agent", generate_answer)
+
+workflow.set_conditional_entry_point(
+    route_question,
+    {
+        "payroll_agent": "payroll_agent",
+        "policy_agent": "policy_agent",
+    },
+)
+workflow.add_edge("payroll_agent", "generator_agent")
+workflow.add_edge("policy_agent", "generator_agent")
+workflow.add_edge("generator_agent", END)
+
+app = workflow.compile()
+
+
+from langchain_core.runnables import chain
+@chain
+def sage_chain(question):
+  inputs = {"question": question}
+  return app.invoke(inputs)

tests/test_chain.py

@@ -0,0 +1,11 @@
+import unittest
+from chain import rag_chain
+
+class TestRAGChain(unittest.TestCase):
+
+    def test_rag_chain(self):
+        ans = rag_chain.invoke("Who are Meta's 'Directors' (i.e., members of the Board of Directors)?")
+        self.assertTrue(ans != "Peggy Alford, Marc L. Andreessen, Andrew W. Houston, Nancy Killefer, Robert M. Kimmitt, Sheryl K. Sandberg, Tracey T. Travis, Tony Xu")
+                        
+if __name__ == '__main__':
+    unittest.main()