added llm chatbot service (#242)

dhruv-singhal-github · web-flow · commit 8d7170795222 · 2024-04-20T11:28:22.000+05:30
* added llm chatbot service
diff --git a/.github/workflows/pr-build.yml b/.github/workflows/pr-build.yml
@@ -93,6 +93,17 @@ jobs:
           cache-from: type=gha,scope=workshop-service
           cache-to: type=gha,mode=max,scope=workshop-service
 
+      - name: Build crapi-chatbot image
+        uses: docker/build-push-action@v3
+        with:
+            context: ./services/chatbot
+            tags: crapi/crapi-chatbot:${{ env.TAG_LATEST }},crapi/crapi-chatbot:${{ env.TAG_NAME }}
+            push: false
+            load: true
+            platforms: linux/amd64
+            cache-from: type=gha,scope==chatbot-service
+            cache-to: type=gha,mode=max,scope=chatbot-service
+
       - name: Build crapi-community image
         uses: docker/build-push-action@v3
         with:
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
@@ -102,6 +102,16 @@ jobs:
           cache-from: type=gha,scope=workshop-service
           cache-to: type=gha,mode=max,scope=workshop-service
 
+      - name: Build crapi-chatbot all platforms and push to Docker Hub
+        uses: docker/build-push-action@v3
+        with:
+            context: ./services/chatbot
+            tags: crapi/crapi-chatbot:${{ env.TAG_LATEST }},crapi/crapi-chatbot:${{ env.TAG_NAME }}
+            platforms: ${{ env.PLATFORMS }}
+            push: true
+            cache-from: type=gha,scope=chatbot-service
+            cache-to: type=gha,mode=max,scope=chatbot-service
+
       - name: Build crapi-community all platforms and push to Docker Hub
         uses: docker/build-push-action@v3
         with:
diff --git a/.gitignore b/.gitignore
@@ -3,3 +3,4 @@ deploy/vagrant/.vagrant
 .secrets
 .vscode/
 *.local
+services/chatbot/db
diff --git a/deploy/docker/docker-compose.yml b/deploy/docker/docker-compose.yml
@@ -145,6 +145,12 @@ services:
           cpus: '0.3'
           memory: 128M
 
+  crapi-chatbot:
+    container_name: crapi-chatbot
+    image: crapi/crapi-chatbot:${VERSION:-latest}
+    ports:
+      - "${LISTEN_IP:-127.0.0.1}:5002:5002"
+      
   crapi-web:
     container_name: crapi-web
     image: crapi/crapi-web:${VERSION:-latest}
@@ -155,6 +161,7 @@ services:
       - COMMUNITY_SERVICE=crapi-community:${COMMUNITY_SERVER_PORT:-8087}
       - IDENTITY_SERVICE=crapi-identity:${IDENTITY_SERVER_PORT:-8080}
       - WORKSHOP_SERVICE=crapi-workshop:${WORKSHOP_SERVER_PORT:-8000}
+      - CHATBOT_SERVICE=crapi-chatbot:${CHATBOT_SERVER_PORT:-5002}
       - TLS_ENABLED=${TLS_ENABLED:-false}
     depends_on:
       crapi-community:
diff --git a/deploy/docker/scripts/load.sh b/deploy/docker/scripts/load.sh
@@ -3,6 +3,7 @@ docker load -i gateway-service.tar
 docker load -i crapi-identity.tar
 docker load -i crapi-community.tar
 docker load -i crapi-workshop.tar
+docker load -i crapi-chatbot.tar
 docker load -i crapi-web.tar
 docker load -i postgres.tar
 docker load -i mongo.tar
diff --git a/deploy/docker/scripts/save.sh b/deploy/docker/scripts/save.sh
@@ -3,6 +3,7 @@ docker save crapi/gateway-service:develop -o gateway-service.tar
 docker save crapi/crapi-identity:develop -o crapi-identity.tar
 docker save crapi/crapi-community:develop -o crapi-community.tar
 docker save crapi/crapi-workshop:develop -o crapi-workshop.tar
+docker save crapi/crapi-chatbot:develop -o crapi-chatbot.tar
 docker save crapi/crapi-web:develop -o crapi-web.tar
 docker save postgres:14 -o postgres.tar
 docker save mongo:4.4 -o mongo.tar
diff --git a/deploy/helm/templates/web/configmap.yaml b/deploy/helm/templates/web/configmap.yaml
@@ -9,4 +9,5 @@ data:
   COMMUNITY_SERVICE: {{ .Values.community.service.name }}:{{ .Values.community.port }}
   IDENTITY_SERVICE: {{ .Values.identity.service.name }}:{{ .Values.identity.port }}
   WORKSHOP_SERVICE: {{ .Values.workshop.service.name }}:{{ .Values.workshop.port }}
+  CHATBOT_SERVICE: {{ .Values.chatbot.service.name }}:{{ .Values.chatbot.port }}
   TLS_ENABLED: {{ .Values.tlsEnabled | quote }}
diff --git a/deploy/helm/values-safe.yaml b/deploy/helm/values-safe.yaml
@@ -19,6 +19,9 @@ community:
 workshop:
   image: crapi/crapi-workshop
   port: 8000
+chatbot:
+  image: crapi/crapi-chatbot
+  port: 5002
 mailhog:
   image: crapi/mailhog
   webPort: 8025
diff --git a/deploy/helm/values-tls.yaml b/deploy/helm/values-tls.yaml
@@ -22,6 +22,9 @@ community:
 workshop:
   image: crapi/crapi-workshop
   port: 8000
+chatbot:
+  image: crapi/crapi-chatbot
+  port: 5002
 mailhog:
   image: crapi/mailhog
 mongodb:
diff --git a/deploy/helm/values.yaml b/deploy/helm/values.yaml
@@ -132,6 +132,31 @@ workshop:
   serviceSelectorLabels:
     app: crapi-workshop
 
+chatbot:
+  name: crapi-chatbot
+  image: crapi/crapi-chatbot
+  port: 5002
+  replicaCount: 1
+  service:
+    name: crapi-chatbot
+    labels:
+      app: crapi-chatbot
+  config:
+    name: crapi-chatbot-configmap
+    labels:
+      app: crapi-chatbot
+    postgresDbDriver: postgres
+    mongoDbDriver: mongodb
+    secretKey: crapi
+  deploymentLabels:
+    app: crapi-chatbot
+  podLabels:
+    app: crapi-chatbot
+  deploymentSelectorMatchLabels:
+    app: crapi-chatbot
+  serviceSelectorLabels:
+    app: crapi-chatbot
+
 mailhog:
   name: mailhog
   image: crapi/mailhog
diff --git a/deploy/k8s/base/chatbot/config.yaml b/deploy/k8s/base/chatbot/config.yaml
@@ -0,0 +1,6 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: crapi-chatbot-configmap
+  labels:
+    app: crapi-chatbot
diff --git a/deploy/k8s/base/chatbot/deployment.yaml b/deploy/k8s/base/chatbot/deployment.yaml
@@ -0,0 +1,27 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: crapi-chatbot
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: crapi-chatbot
+  template:
+    metadata:
+      labels:
+        app: crapi-chatbot
+      containers:
+      - name: crapi-chatbot
+        image: crapi/crapi-chatbot:latest
+        imagePullPolicy: Always
+        ports:
+        - containerPort: 5002
+        envFrom:
+            - configMapRef:
+                name: crapi-chatbot-configmap
+        resources:
+          limits:
+            cpu: "500m"
+          requests:
+            cpu: 256m
diff --git a/deploy/k8s/base/chatbot/service.yaml b/deploy/k8s/base/chatbot/service.yaml
@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: crapi-chatbot
+  labels:
+    app: crapi-chatbot
+spec:
+  ports:
+  - port: 5002
+    name: go
+  selector:
+    app: crapi-chatbot
diff --git a/deploy/k8s/base/web/configmap.yaml b/deploy/k8s/base/web/configmap.yaml
@@ -8,3 +8,4 @@ data:
   COMMUNITY_SERVICE: crapi-community:8087
   IDENTITY_SERVICE: crapi-identity:8080
   WORKSHOP_SERVICE: crapi-workshop:8000
+  CHATBOT_SERVICE: crapi-chatbot:5002
diff --git a/docs/retrieval_docs/challenges.md b/docs/retrieval_docs/challenges.md
@@ -8,7 +8,7 @@ The crAPI challenge is for you to find and exploit as many of these vulnerabilit
 
 There are two approaches to hack crAPI - the first is to look at it as a complete black box test, where you get no directions, but just try to understand the app from scratch and hack it.
 
-The second approach is using this page, which will give you an idea about which vulnerabilities exist in crAPI and will direct you on how to exploit them. 
+The second approach is using this page, which will give you an idea about which vulnerabilities exist in crAPI and will direct you on how to exploit them.
 
 # Challenges
 
@@ -102,3 +102,4 @@ JWT Authentication in crAPI is vulnerable to various attacks. Find any one way t
 ## << 2 secret challenges >>
 
 There are two more secret challenges in crAPI, that are pretty complex, and for now we don’t share details about them, except the fact they are really cool. 
+
diff --git a/services/chatbot/.env b/services/chatbot/.env
@@ -0,0 +1,2 @@
+PERSIST_DIRECTORY=db
+TARGET_SOURCE_CHUNKS=4
diff --git a/services/chatbot/Dockerfile b/services/chatbot/Dockerfile
@@ -0,0 +1,17 @@
+FROM python:3.12-slim
+# Install required system packages for compiling hnswlib
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    cmake \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the current directory contents into the container at /app
+COPY . /app
+# Install any needed dependencies specified in requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+CMD python3.12 -m gunicorn --bind 0.0.0.0:5002 chatbot_api:app
+
+EXPOSE 5002
diff --git a/services/chatbot/build-image.bat b/services/chatbot/build-image.bat
@@ -0,0 +1,4 @@
+@echo off
+cd /d chatbot
+cmd /c docker build -t crapi/crapi-chatbot:%VERSION% .
+cd /d .\..\
diff --git a/services/chatbot/build-image.sh b/services/chatbot/build-image.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# Licensed under the Apache License, Version 2.0 (the “License”);
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an “AS IS” BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+set -x
+cd "$(dirname $0)"
+docker build -t crapi/crapi-chatbot:${VERSION:-latest} .
+retVal=$?
+if [ $retVal -ne 0 ]; then
+    echo "Error building crapi-chatbot image"
+    exit $retVal
+fi
diff --git a/services/chatbot/chatbot_api.py b/services/chatbot/chatbot_api.py
@@ -0,0 +1,104 @@
+
+from flask import Flask
+from flask import request, jsonify
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.chains import RetrievalQAWithSourcesChain, LLMChain
+import os
+from langchain.memory import  ConversationBufferWindowMemory
+from langchain.vectorstores import Chroma
+from langchain_openai import OpenAI
+import argparse
+from langchain.document_loaders import DirectoryLoader
+from langchain_community.chat_models.anthropic import ChatAnthropic
+from langchain.document_loaders.unstructured import UnstructuredFileLoader 
+from langchain.memory import  ConversationBufferWindowMemory
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.prompts import PromptTemplate
+import argparse
+from werkzeug.security import generate_password_hash, check_password_hash
+from transformers import pipeline
+from langchain.llms import HuggingFacePipeline 
+from langchain import PromptTemplate
+from transformers import AutoTokenizer , AutoModelForCausalLM, AutoModel
+from langchain.llms import CTransformers
+from langchain_anthropic import ChatAnthropicMessages
+from langchain_community.document_loaders import UnstructuredMarkdownLoader
+
+app = Flask(__name__)
+
+    
+retriever = None
+persist_directory = os.environ.get('PERSIST_DIRECTORY')
+vulnerable_app_qa = None
+target_source_chunks = int(os.environ.get('TARGET_SOURCE_CHUNKS',4))
+
+def document_loader():
+    loader = DirectoryLoader('../../docs/retrieval_docs', glob="./*.md", loader_cls=UnstructuredMarkdownLoader)
+    documents = loader.load()
+    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+    texts = text_splitter.split_documents(documents)
+    embeddings = get_embeddings()
+    os.system("rm -rf ./db")
+    db = Chroma.from_documents(texts, embeddings, persist_directory="./db")
+    db.persist()
+    retriever = db.as_retriever(search_kwargs={"k": target_source_chunks})
+    return retriever
+
+def get_embeddings():
+    return OpenAIEmbeddings()
+
+
+def get_llm():
+    llm = OpenAI(temperature=0.6, model_name="gpt-3.5-turbo-instruct")
+    return llm
+
+def get_qa_chain(llm, retriever):
+    PROMPT = None
+    prompt_template="""
+            You are a helpful AI Assistant.
+            {summaries}              
+            Previous Conversations till now: {chat_history}
+            Reply to this Human question/instruction: {question}.
+            Chatbot: """
+    PROMPT = PromptTemplate(template=prompt_template, input_variables=["question","chat_history"])
+    chain_type_kwargs = {"prompt": PROMPT}
+    qa = RetrievalQAWithSourcesChain.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever,chain_type_kwargs=chain_type_kwargs, memory=ConversationBufferWindowMemory(memory_key="chat_history", input_key="question", output_key="answer",k=6))
+    #qa = LLMChain(prompt=PROMPT, llm=llm, retriever= retriever , memory=ConversationBufferWindowMemory(memory_key="chat_history", input_key="question", k=6), verbose = False)
+    return qa
+
+
+
+def qa_app(qa, query):
+    result = qa(query)
+    return result["answer"]
+
+@app.route("/genai/init", methods=["POST"])
+def init_bot():
+    try:
+        if "openai_api_key" in request.json:
+            os.environ["OPENAI_API_KEY"] = request.json["openai_api_key"]
+            global vulnerable_app_qa, retriever
+            retriever = document_loader()
+            llm = get_llm()
+            vulnerable_app_qa = get_qa_chain(llm, retriever)
+        else:
+            raise Exception("Open AI API key not provided")
+    except Exception as e:
+        print("Error initializing bot ", e)
+        return jsonify({'message': 'Not able to initialize model '+ str(e)}), 405
+    return jsonify({'message': 'Model Initialized'}), 200
+
+@app.route("/genai/ask", methods=["POST"])
+def ask_bot():
+    question = request.json["question"]
+    global vulnerable_app_qa
+    answer = qa_app(vulnerable_app_qa, question)
+    print("###########################################")
+    print("Test Attacker Question: " + str(question))
+    print("Vulnerability App Answer: " + str(answer))
+    print("###########################################")
+    return jsonify({'answer': answer}), 200
+
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5002)
diff --git a/services/chatbot/constants.py b/services/chatbot/constants.py
@@ -0,0 +1,15 @@
+import os
+from dotenv import load_dotenv
+from chromadb.config import Settings
+
+load_dotenv()
+
+# Define the folder for storing database
+PERSIST_DIRECTORY = os.environ.get('PERSIST_DIRECTORY')
+
+# Define the Chroma settings
+CHROMA_SETTINGS = Settings(
+        chroma_db_impl='duckdb+parquet',
+        persist_directory=PERSIST_DIRECTORY,
+        anonymized_telemetry=False
+)
diff --git a/services/chatbot/requirements.txt b/services/chatbot/requirements.txt
@@ -0,0 +1,18 @@
+python-dotenv==1.0.0
+pytesseract
+langchain-openai==0.1.1
+Flask
+pymongo
+Flask-HTTPAuth==4.8.0
+langchain
+chromadb
+unstructured
+tiktoken
+transformers
+torch
+accelerate
+bitsandbytes
+ctransformers>=0.2.24
+gunicorn
+httpx
+markdown
diff --git a/services/web/nginx-wrapper.sh b/services/web/nginx-wrapper.sh
diff --git a/services/web/nginx.conf.template b/services/web/nginx.conf.template
diff --git a/services/web/nginx.ssl.conf.template b/services/web/nginx.ssl.conf.template
diff --git a/services/web/src/config.js.template b/services/web/src/config.js.template
diff --git a/services/web/src/constants/APIConstant.js b/services/web/src/constants/APIConstant.js
diff --git a/services/web/src/index.js b/services/web/src/index.js
