Upload 24 files

.gitignore

@@ -0,0 +1,165 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Win executables
+*.exe
+rag-env/
+mixtral-playground/
+
+# 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,22 @@
+FROM python:3.9-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    software-properties-common \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY components /app/
+COPY middlewares /app/
+COPY app.py /app/
+COPY requirements.txt /app/
+COPY config.yaml /app/
+
+RUN pip3 install -r requirements.txt
+
+EXPOSE 8501
+
+ENTRYPOINT ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]

README.md

@@ -1,12 +1,41 @@
 ---
-title: NCTC OSINT AGENT
-emoji: 🌖
-colorFrom: indigo
-colorTo: green
+title: Mixtral Search Engine
+emoji: 🔍
+colorFrom: pink
+colorTo: gray
 sdk: streamlit
-sdk_version: 1.33.0
+sdk_version: 1.29.0
 app_file: app.py
 pinned: false
+license: mit
 ---
 
+# Mixtral Search Engine
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## Docker Setup
+
+If you prefer using Docker, follow these steps:
+
+1. Clone the repository.
+2. Create a `.env` file to store API credentials.
+
+   ```
+   HF_TOKEN = ...
+   GOOGLE_SEARCH_ENGINE_ID = ...
+   GOOGLE_SEARCH_API_KEY = ...
+   BING_SEARCH_API_KEY = ...
+   ```
+
+3. Build docker image using
+
+   ```
+   docker build -t mixtral-search .
+   ```
+
+4. Run the image using
+
+   ```
+   docker run --env-file .env -p 8501:8501 mixtral-search
+   ```

app.py

@@ -0,0 +1,32 @@
+import yaml
+import streamlit as st
+from components.sidebar import sidebar
+from components.chat_box import chat_box
+from components.chat_loop import chat_loop
+from components.init_state import init_state
+from components.prompt_engineering_dashboard import prompt_engineering_dashboard
+
+
+
+with open("config.yaml", "r") as file:
+    config = yaml.safe_load(file)
+
+st.set_page_config(
+    page_title="NCTC OSINT AGENT",
+    page_icon="📚",
+)
+
+
+init_state(st.session_state, config)
+
+st.write("# NCTC OSINT AGENT ")
+
+# Prompt Engineering Dashboard is working but not for production, works great for testing.
+prompt_engineering_dashboard(st.session_state, config)
+
+
+sidebar(st.session_state, config)
+
+chat_box(st.session_state, config)
+
+chat_loop(st.session_state, config)

components/chat_box.py

@@ -0,0 +1,7 @@
+import streamlit as st
+
+
+def chat_box(session_state, config):
+    for message in session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])

components/chat_loop.py

@@ -0,0 +1,23 @@
+import streamlit as st
+from components.generate_chat_stream import generate_chat_stream
+from components.stream_handler import stream_handler
+from components.show_source import show_source
+
+
+def chat_loop(session_state, config):
+    if prompt := st.chat_input("Search the web..."):
+        st.chat_message("user").markdown(prompt)
+        session_state.messages.append({"role": "user", "content": prompt})
+
+        chat_stream, links = generate_chat_stream(session_state, prompt, config)
+
+        with st.chat_message("assistant"):
+            placeholder = st.empty()
+            full_response = stream_handler(
+                session_state, chat_stream, prompt, placeholder
+            )
+            if session_state.rag_enabled:
+                show_source(links)
+
+        session_state.history.append([prompt, full_response])
+        session_state.messages.append({"role": "assistant", "content": full_response})

components/generate_chat_stream.py

@@ -0,0 +1,104 @@
+import streamlit as st
+from middlewares.utils import gen_augmented_prompt_via_websearch
+from middlewares.chat_client import chat
+import json
+from pprint import  pformat
+from notion_client import Client
+
+def safe_get(data, dot_chained_keys):
+    '''
+        {'a': {'b': [{'c': 1}]}}
+        safe_get(data, 'a.b.0.c') -> 1
+    '''
+    keys = dot_chained_keys.split('.')
+    for key in keys:
+        try:
+            if isinstance(data, list):
+                data = data[int(key)]
+            else:
+                data = data[key]
+        except (KeyError, TypeError, IndexError):
+            return None
+    return data
+
+def get_notion_data() :
+    integration_token = "secret_lTOe0q9dqqKQLRRb2KJwi7QFSl0vqoztroRFHW6MeQE"
+    notion_database_id = "6c0d877b823a4e3699016fa7083f3006"
+
+    client = Client(auth=integration_token)
+
+    first_db_rows = client.databases.query(notion_database_id)
+    rows = []
+
+
+    for row in first_db_rows['results']:
+        price = safe_get(row, 'properties.($) Per Unit.number')
+        store_link = safe_get(row,  'properties.Store Link.url')
+        supplier_email = safe_get(row,  'properties.Supplier Email.email')
+        exp_del = safe_get(row,  'properties.Expected Delivery.date')
+        
+        collections = safe_get(row,  'properties.Collection.multi_select')
+        collection_names = []
+        for collection in collections :
+            collection_names.append(collection['name'])
+
+        status = safe_get(row,  'properties.Status.select.name')
+        sup_phone = safe_get(row, 'properties.Supplier Phone.phone_number')
+        stock_alert = safe_get(row, 'properties.Stock Alert.status.name') 
+        prod_name = safe_get(row, 'properties.Product .title.0.text.content')
+        sku = safe_get(row, 'properties.SKU.number')
+        shipped_date = safe_get(row, 'properties.Shipped On.date')
+        on_order = safe_get(row, 'properties.On Order.number')
+        on_hand = safe_get(row, 'properties.On Hand.number')
+        size_names = []
+        sizes = safe_get(row, 'properties.Size.multi_select')
+        for size in sizes :
+            size_names.append(size['name'])
+        
+        rows.append({
+        'Price Per unit': price,
+        'Store Link' : store_link,
+        'Supplier Email' : supplier_email,
+        'Expected Delivery' : exp_del,
+        'Collection' : collection_names,
+        'Status' : status,
+        'Supplier Phone' : sup_phone,
+        'Stock Alert' : stock_alert,
+        'Product Name' : prod_name,
+        'SKU' : sku,
+        'Sizes' : size_names,
+        'Shipped Date' : shipped_date, 
+        'On Order' : on_order,
+        "On Hand" : on_hand,
+        })
+
+    notion_data_string = pformat(rows)
+    return notion_data_string
+
+def generate_chat_stream(session_state, query, config):
+    # 1. augments prompt according to the template
+    # 2. returns chat_stream and source links
+    # 3. chat_stream and source links are used by stream_handler and show_source
+    chat_bot_dict = config["CHAT_BOTS"]
+    links = []
+    if session_state.rag_enabled:
+        with st.spinner("Fetching relevent documents from Web...."):
+            query, links = gen_augmented_prompt_via_websearch(
+                prompt=query,
+                pre_context=session_state.pre_context,
+                post_context=session_state.post_context,
+                pre_prompt=session_state.pre_prompt,
+                post_prompt=session_state.post_prompt,
+                search_vendor=session_state.search_vendor,
+                top_k=session_state.top_k,
+                n_crawl=session_state.n_crawl,
+                pass_prev=session_state.pass_prev,
+                prev_output=session_state.history[-1][1],
+            )
+    
+    notion_data = get_notion_data()
+
+    with st.spinner("Generating response..."):
+        chat_stream = chat(session_state, notion_data + " " + query , config)
+
+    return chat_stream, links

components/init_state.py

@@ -0,0 +1,64 @@
+def init_state(session_state, config):
+    initial_prompt_engineering_dict = config["PROMPT_ENGINEERING_DICT"]
+    if "messages" not in session_state:
+        session_state.messages = []
+
+    if "tokens_used" not in session_state:
+        session_state.tokens_used = 0
+
+    if "tps" not in session_state:
+        session_state.tps = 0
+
+    if "temp" not in session_state:
+        session_state.temp = 0.8
+
+    if "history" not in session_state:
+        session_state.history = [
+            [
+                initial_prompt_engineering_dict["SYSTEM_INSTRUCTION"],
+                initial_prompt_engineering_dict["SYSTEM_RESPONSE"],
+            ]
+        ]
+
+    if "n_crawl" not in session_state:
+        session_state.n_crawl = 5
+
+    if "repetion_penalty" not in session_state:
+        session_state.repetion_penalty = 1
+
+    if "rag_enabled" not in session_state:
+        session_state.rag_enabled = True
+
+    if "chat_bot" not in session_state:
+        session_state.chat_bot = "Mixtral 8x7B v0.1"
+
+    if "search_vendor" not in session_state:
+        session_state.search_vendor = "Bing"
+
+    if "system_instruction" not in session_state:
+        session_state.system_instruction = initial_prompt_engineering_dict[
+            "SYSTEM_INSTRUCTION"
+        ]
+
+    if "system_response" not in session_state:
+        session_state.system_instruction = initial_prompt_engineering_dict[
+            "SYSTEM_RESPONSE"
+        ]
+
+    if "pre_context" not in session_state:
+        session_state.pre_context = initial_prompt_engineering_dict["PRE_CONTEXT"]
+
+    if "post_context" not in session_state:
+        session_state.post_context = initial_prompt_engineering_dict["POST_CONTEXT"]
+
+    if "pre_prompt" not in session_state:
+        session_state.pre_prompt = initial_prompt_engineering_dict["PRE_PROMPT"]
+
+    if "post_prompt" not in session_state:
+        session_state.post_prompt = initial_prompt_engineering_dict["POST_PROMPT"]
+
+    if "pass_prev" not in session_state:
+        session_state.pass_prev = False
+
+    if "chunk_size" not in session_state:
+        session_state.chunk_size = 512

components/prompt_engineering_dashboard.py

@@ -0,0 +1,68 @@
+import streamlit as st
+
+
+def prompt_engineering_dashboard(session_state, config):
+    inital_prompt_engineering_dict = config["PROMPT_ENGINEERING_DICT"]
+
+    def engineer_prompt():
+        session_state.history[0] = [
+            session_state.system_instruction,
+            session_state.system_response,
+        ]
+
+    with st.expander("Prompt Engineering Dashboard"):
+        st.info(
+            "**The input to the model follows this below template**",
+        )
+        st.code(
+            """
+                    [SYSTEM INSTRUCTION]
+                    [SYSTEM RESPONSE]
+                    
+                    [... LIST OF PREV INPUTS]
+                    
+                    [PRE CONTEXT]
+                    [CONTEXT RETRIEVED FROM THE WEB]
+                    [POST CONTEXT]
+
+                    [PRE PROMPT]
+                    [PROMPT]
+                    [POST PROMPT]
+                    [PREV GENERATED INPUT] # Only if  Pass previous prompt set True  
+
+                    """
+        )
+        session_state.system_instruction = st.text_area(
+            label="SYSTEM INSTRUCTION",
+            value=inital_prompt_engineering_dict["SYSTEM_INSTRUCTION"],
+        )
+        session_state.system_response = st.text_area(
+            "SYSTEM RESPONSE", value=inital_prompt_engineering_dict["SYSTEM_RESPONSE"]
+        )
+
+        col1, col2 = st.columns(2)
+        with col1:
+            session_state.pre_context = st.text_input(
+                "PRE CONTEXT",
+                value=inital_prompt_engineering_dict["PRE_CONTEXT"],
+                disabled=not session_state.rag_enabled,
+            )
+            session_state.post_context = st.text_input(
+                "POST CONTEXT",
+                value=inital_prompt_engineering_dict["POST_CONTEXT"],
+                disabled=not session_state.rag_enabled,
+            )
+
+        with col2:
+            session_state.pre_prompt = st.text_input(
+                "PRE PROMPT", value=inital_prompt_engineering_dict["PRE_PROMPT"]
+            )
+            session_state.post_prompt = st.text_input(
+                "POST PROMPT", value=inital_prompt_engineering_dict["POST_PROMPT"]
+            )
+
+        col3, col4 = st.columns(2)
+        with col3:
+            session_state.pass_prev = st.toggle("Pass previous Output")
+        with col4:
+            st.button("Engineer Prompts", on_click=engineer_prompt)

components/show_source.py

@@ -0,0 +1,8 @@
+import streamlit as st
+
+
+def show_source(links):
+    # Expander component to show source
+    with st.expander("Show source"):
+        for i, link in enumerate(links):
+            st.info(f"{link}")

components/sidebar.py

@@ -0,0 +1,11 @@
+import streamlit as st
+from components.sidebar_components.model_analytics import model_analytics
+from components.sidebar_components.retrieval_settings import retrieval_settings
+from components.sidebar_components.model_settings  import model_settings
+
+
+def sidebar(session_state, config):
+    with st.sidebar:
+        retrieval_settings(session_state, config)
+        model_analytics(session_state, config)
+        model_settings(session_state, config)

components/sidebar_components/model_analytics.py

@@ -0,0 +1,18 @@
+import streamlit as st
+def model_analytics(session_state, config):
+    COST_PER_1000_TOKENS_USD = config["COST_PER_1000_TOKENS_USD"]
+    
+    st.markdown("# Model Analytics")
+
+    st.write("Total tokens used :", session_state["tokens_used"])
+    st.write("Speed :", session_state["tps"], "  tokens/sec")
+    st.write(
+        "Total cost incurred :",
+        round(
+            COST_PER_1000_TOKENS_USD * session_state["tokens_used"] / 1000,
+            3,
+        ),
+        "USD",
+    )
+
+    st.markdown("---")

components/sidebar_components/model_settings.py

@@ -0,0 +1,25 @@
+import streamlit as st
+
+def model_settings(session_state,config):
+    CHAT_BOTS = config["CHAT_BOTS"]
+    
+    st.markdown("# Model Settings")
+
+    session_state.chat_bot = st.sidebar.radio(
+        "Select one:", [key for key, _ in CHAT_BOTS.items()]
+    )
+    session_state.temp = st.slider(
+        label="Temperature", min_value=0.0, max_value=1.0, step=0.1, value=0.9
+    )
+
+    session_state.max_tokens = st.slider(
+        label="New tokens to generate",
+        min_value=64,
+        max_value=2048,
+        step=32,
+        value=512,
+    )
+
+    session_state.repetion_penalty = st.slider(
+        label="Repetion Penalty", min_value=0.0, max_value=1.0, step=0.1, value=1.0
+    )

components/sidebar_components/retrieval_settings.py

@@ -0,0 +1,37 @@
+import streamlit as st
+
+def retrieval_settings(session_state, config):
+        st.markdown("# Web Retrieval")
+        session_state.rag_enabled = st.toggle("Activate Web Retrieval", value=True)
+        session_state.search_vendor = st.radio(
+            "Select Search Vendor",
+            ["Bing", "Google"],
+            disabled=not session_state.rag_enabled,
+        )
+        session_state.n_crawl = st.slider(
+            label="Links to Crawl",
+            key=1,
+            min_value=1,
+            max_value=10,
+            value=4,
+            disabled=not session_state.rag_enabled,
+        )
+        session_state.top_k = st.slider(
+            label="Chunks to Retrieve via Reranker",
+            key=2,
+            min_value=1,
+            max_value=20,
+            value=5,
+            disabled=not session_state.rag_enabled,
+        )
+
+        session_state.chunk_size = st.slider(
+            label="Chunk Size",
+            value=512,
+            min_value=128,
+            max_value=1024,
+            step=8,
+            disabled=not session_state.rag_enabled,
+        )
+
+        st.markdown("---")

components/stream_handler.py

@@ -0,0 +1,41 @@
+import time
+import streamlit as st
+
+COST_PER_1000_TOKENS_USD = 0.139 / 80
+
+
+def stream_handler(session_state, chat_stream, prompt, placeholder):
+    # 1. Uses the chat_stream and streams message on placeholder
+    # 2. returns full_response for token calculation
+    start_time = time.time()
+    full_response = ""
+
+    for chunk in chat_stream:
+        if chunk.token.text in ["</s>", "<|im_end|>"]:
+            break;
+        full_response += chunk.token.text
+        placeholder.markdown(full_response + "▌")
+    placeholder.markdown(full_response)
+
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    total_tokens_processed = len(full_response.split())
+    tokens_per_second = total_tokens_processed // elapsed_time
+    len_response = (len(prompt.split()) + len(full_response.split())) * 1.25
+    col1, col2, col3 = st.columns(3)
+
+    with col1:
+        st.write(f"**{tokens_per_second} tokens/second**")
+
+    with col2:
+        st.write(f"**{int(len_response)} tokens generated**")
+
+    with col3:
+        st.write(
+            f"**$ {round(len_response * COST_PER_1000_TOKENS_USD  / 1000, 5)} cost incurred**"
+        )
+
+    session_state["tps"] = tokens_per_second
+    session_state["tokens_used"] = len_response + session_state["tokens_used"]
+
+    return full_response

config.yaml

@@ -0,0 +1,24 @@
+PROMPT_ENGINEERING_DICT:
+  SYSTEM_INSTRUCTION: |
+    YOU ARE A SEARCH ENGINE AND AN INVENTORY MANAGER HAVING FULL ACCESS TO WEB PAGES AND NOTION DATABASE IN JSON, 
+    YOU GIVE EXTREMELY DETAILED AND ACCURATE INFORMATION ACCORDING TO USER PROMPTS.
+
+  SYSTEM_RESPONSE: |
+    Certainly! I'm here to help. What information are you looking for? 
+    Please provide me with a specific topic or question, and I'll do my 
+    best to provide you with detailed and accurate information.
+
+  PRE_CONTEXT: NOW YOU ARE SEARCHING THE WEB, AND HERE ARE THE CHUNKS RETRIEVED FROM THE WEB, YOU ALSO HAVE ACCESS TO INVENTORY DATASET IN JSON FORMAT.
+  POST_CONTEXT: ""
+  PRE_PROMPT: NOW ACCORDING TO THE CONTEXT RETRIEVED FROM THE GENERATE THE CONTENT FOR THE FOLLOWING SUBJECT
+  POST_PROMPT: PRIORITIZE DATA, FACTS AND STATISTICS OVER PERSONAL EXPERIENCES AND OPINIONS, FOCUS MORE ON STATISTICS AND DATA.
+
+CHAT_BOTS:
+  Nous Hermes 2: NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO
+  Mixtral 8x7B v0.1: mistralai/Mixtral-8x7B-Instruct-v0.1
+  Mistral 7B v0.1: mistralai/Mistral-7B-Instruct-v0.1
+  Mistral 7B v0.2: mistralai/Mistral-7B-Instruct-v0.2
+
+CROSS_ENCODERS:
+
+COST_PER_1000_TOKENS_USD: 0.001737375

middlewares/chat_client.py

@@ -0,0 +1,78 @@
+from huggingface_hub import InferenceClient
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+API_TOKEN = os.getenv("HF_TOKEN")
+
+
+
+
+def format_prompt(session_state ,query, history, chat_client):
+    if chat_client=="NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO" :
+        model_input = f"""<|im_start|>system
+        {session_state.system_instruction}
+        """
+        for user_prompt, bot_response in history:
+            model_input += f"""<|im_start|>user
+            {user_prompt}<|im_end|>
+            """
+            model_input += f"""<|im_start|>assistant
+            {bot_response}<|im_end|>
+            """
+        model_input += f"""<|im_start|>user 
+        {query}<|im_end|>
+        <|im_start|>assistant"""
+        
+        return model_input
+
+        
+    else :
+        model_input = "<s>"
+        for user_prompt, bot_response in history:
+            model_input += f"[INST] {user_prompt} [/INST]"
+            model_input += f" {bot_response}</s> "
+        model_input += f"[INST] {query} [/INST]"
+        return model_input
+
+
+def chat(session_state, query, config):
+    
+
+    
+    chat_bot_dict = config["CHAT_BOTS"]
+    chat_client = chat_bot_dict[session_state.chat_bot]
+    temperature = session_state.temp
+    max_new_tokens = session_state.max_tokens
+    repetion_penalty = session_state.repetion_penalty
+    history = session_state.history
+
+
+    client = InferenceClient(chat_client, token=API_TOKEN)
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(0.95)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetion_penalty,
+        do_sample=True,
+        seed=42,
+    )
+
+    formatted_prompt = format_prompt(session_state, query, history, chat_client)
+
+    stream = client.text_generation(
+        formatted_prompt,
+        **generate_kwargs,
+        stream=True,
+        details=True,
+        return_full_text=False,
+        truncate = 32000
+    )
+
+    return stream

middlewares/search_client.py

@@ -0,0 +1,85 @@
+import requests
+from bs4 import BeautifulSoup
+import re
+import concurrent.futures
+
+
+class SearchClient:
+    def __init__(self, vendor, engine_id=None, api_key=None):
+        self.vendor = vendor
+        if vendor == "google":
+            self.endpoint = f"https://www.googleapis.com/customsearch/v1?key={api_key}&cx={engine_id}"
+        elif vendor == "bing":
+            self.endpoint = "https://api.bing.microsoft.com/v7.0/search"
+            self.headers = {
+                "Ocp-Apim-Subscription-Key": api_key,
+            }
+
+    @staticmethod
+    def _extract_text_from_link(link):
+        page = requests.get(link)
+        if page.status_code == 200:
+            soup = BeautifulSoup(page.content, "html.parser")
+            text = soup.get_text()
+            cleaned_text = re.sub(r"\s+", " ", text)
+            return cleaned_text
+        return None
+
+    def _fetch_text_from_links(self, links):
+        results = []
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future_to_link = {
+                executor.submit(self._extract_text_from_link, link): link
+                for link in links
+            }
+            for future in concurrent.futures.as_completed(future_to_link):
+                link = future_to_link[future]
+                try:
+                    cleaned_text = future.result()
+                    if cleaned_text:
+                        results.append({"text": cleaned_text, "link": link})
+                except Exception as e:
+                    print(f"Error fetching data from {link}: {e}")
+        return results
+
+    def _google_search(self, query, n_crawl):
+        response = requests.get(self.endpoint, params={"q": query})
+        search_results = response.json()
+
+        results = []
+        count = 0
+        for item in search_results.get("items", []):
+            if count >= n_crawl:
+                break
+
+            link = item["link"]
+            results.append(link)
+            count += 1
+
+        text_results = self._fetch_text_from_links(results)
+        return text_results
+
+    def _bing_search(self, query, n_crawl):
+        params = {
+            "q": query,
+            "count": n_crawl,  # You might need to adjust this based on Bing API requirements
+            "mkt": "en-US",
+        }
+        response = requests.get(self.endpoint, headers=self.headers, params=params)
+        search_results = response.json()
+
+        results = []
+        for item in search_results.get("webPages", {}).get("value", []):
+            link = item["url"]
+            results.append(link)
+
+        text_results = self._fetch_text_from_links(results)
+        return text_results
+
+    def search(self, query, n_crawl):
+        if self.vendor == "google":
+            return self._google_search(query, n_crawl)
+        elif self.vendor == "bing":
+            return self._bing_search(query, n_crawl)
+        else:
+            return "Invalid vendor"

middlewares/utils.py

@@ -0,0 +1,117 @@
+from sentence_transformers import CrossEncoder
+
+import json
+import math
+import numpy as np
+from middlewares.search_client import SearchClient
+import os
+from dotenv import load_dotenv
+
+
+load_dotenv()
+
+
+GOOGLE_SEARCH_ENGINE_ID = os.getenv("GOOGLE_SEARCH_ENGINE_ID")
+GOOGLE_SEARCH_API_KEY = os.getenv("GOOGLE_SEARCH_API_KEY")
+BING_SEARCH_API_KEY = os.getenv("BING_SEARCH_API_KEY")
+
+reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+
+googleSearchClient = SearchClient(
+    "google", api_key=GOOGLE_SEARCH_API_KEY, engine_id=GOOGLE_SEARCH_ENGINE_ID
+)
+bingSearchClient = SearchClient("bing", api_key=BING_SEARCH_API_KEY, engine_id=None)
+
+
+
+
+def rerank(query, top_k, search_results, chunk_size=512):
+    chunks = []
+    for result in search_results:
+        text = result["text"]
+        words = text.split()
+        num_chunks = math.ceil(len(words) / chunk_size)
+        for i in range(num_chunks):
+            start = i * chunk_size
+            end = (i + 1) * chunk_size
+            chunk = " ".join(words[start:end])
+            chunks.append((result["link"], chunk))
+
+    # Create sentence combinations with the query
+    sentence_combinations = [[query, chunk[1]] for chunk in chunks]
+
+    # Compute similarity scores for these combinations
+    similarity_scores = reranker.predict(sentence_combinations)
+
+    # Sort scores indexes in decreasing order
+    sim_scores_argsort = reversed(np.argsort(similarity_scores))
+
+    # Rearrange search_results based on the reranked scores
+    reranked_results = []
+    for idx in sim_scores_argsort:
+        link = chunks[idx][0]
+        chunk = chunks[idx][1]
+        reranked_results.append({"link": link, "text": chunk})
+
+    # Return the top K ranks
+    return reranked_results[:top_k]
+
+
+def gen_augmented_prompt_via_websearch(
+    prompt,
+    search_vendor,
+    n_crawl,
+    top_k,
+    pre_context="",
+    post_context="",
+    pre_prompt="",
+    post_prompt="",
+    pass_prev=False,
+    prev_output="",
+    chunk_size=512,
+):
+   
+
+    search_results = []
+    reranked_results = []
+    if search_vendor == "Google":
+        search_results = googleSearchClient.search(prompt, n_crawl)
+    elif search_vendor == "Bing":
+        search_results = bingSearchClient.search(prompt, n_crawl)
+
+    if len(search_results) > 0:
+        reranked_results = rerank(prompt, top_k, search_results, chunk_size)
+
+    links = []
+    context = ""
+    for res in reranked_results:
+        context += res["text"] + "\n\n"
+        link = res["link"]
+        links.append(link)
+
+    # remove duplicate links
+    links = list(set(links))
+
+    prev_output = prev_output if pass_prev else ""
+
+    augmented_prompt = f"""
+    
+    {pre_context}
+
+    {context}
+
+    
+    {post_context}
+    
+    {pre_prompt} 
+    
+    {prompt} 
+    
+    {post_prompt}
+
+    {prev_output}
+
+    """
+
+    print(augmented_prompt)
+    return augmented_prompt, links

requirements.txt

@@ -0,0 +1,97 @@
+notion-client==2.2.1
+altair==5.1.2
+asttokens==2.2.1
+attrs==23.1.0
+backcall==0.2.0
+beautifulsoup4==4.12.2
+blinker==1.6.3
+cachetools==5.3.1
+certifi==2023.7.22
+charset-normalizer==3.3.0
+click==8.1.7
+colorama==0.4.6
+comm==0.1.3
+debugpy==1.6.7
+decorator==5.1.1
+dnspython==2.4.2
+executing==1.2.0
+filelock==3.12.4
+fsspec==2023.9.2
+gitdb==4.0.10
+GitPython==3.1.37
+huggingface-hub==0.18.0
+idna==3.4
+importlib-metadata==6.8.0
+ipykernel==6.23.3
+ipython==8.14.0
+jedi==0.18.2
+Jinja2==3.1.2
+joblib==1.3.2
+jsonschema==4.19.1
+jsonschema-specifications==2023.7.1
+jupyter_client==8.3.0
+jupyter_core==5.3.1
+loguru==0.7.2
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+mpmath==1.3.0
+nest-asyncio==1.5.6
+networkx==3.2.1
+nltk==3.8.1
+numpy==1.26.0
+packaging==23.1
+pandas==2.1.1
+parso==0.8.3
+pickleshare==0.7.5
+Pillow==10.0.1
+platformdirs==3.8.0
+prompt-toolkit==3.0.38
+protobuf==4.24.4
+psutil==5.9.5
+pure-eval==0.2.2
+pyarrow==13.0.0
+pydeck==0.8.1b0
+Pygments==2.15.1
+python-dateutil==2.8.2
+python-dotenv==1.0.0
+pytz==2023.3.post1
+PyYAML==6.0.1
+pyzmq==25.1.0
+referencing==0.30.2
+regex==2023.10.3
+requests==2.31.0
+rich==13.6.0
+rpds-py==0.10.4
+safetensors==0.4.1
+scikit-learn==1.3.2
+scipy==1.11.4
+sentence-transformers==2.2.2
+sentencepiece==0.1.99
+six==1.16.0
+smmap==5.0.1
+soupsieve==2.5
+stack-data==0.6.2
+streamlit==1.27.2
+sympy==1.12
+tenacity==8.2.3
+threadpoolctl==3.2.0
+tokenizers==0.15.0
+toml==0.10.2
+toolz==0.12.0
+torch==2.1.2
+torchvision==0.16.2
+tornado==6.3.2
+tqdm==4.66.1
+traitlets==5.9.0
+transformers==4.35.2
+typing_extensions==4.8.0
+tzdata==2023.3
+tzlocal==5.1
+urllib3==2.0.6
+validators==0.22.0
+watchdog==3.0.0
+wcwidth==0.2.6
+win32-setctime==1.1.0
+zipp==3.17.0