From 4e1ab8016d5a40f1a5082caed0db7f68dddac6d6 Mon Sep 17 00:00:00 2001
From: Mark Chen <MarkCodering>
Date: Wed, 7 Feb 2024 09:20:49 +0800
Subject: [PATCH] Add GPT-3 and Llama 2 7B chatbot functionality

---
 src/app/main.py   | 44 +++++++++++++++++++++++++++++++++++++-------
 src/index/main.py |  5 +++--
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/src/app/main.py b/src/app/main.py
index 9e118e1..42e2fd2 100644
--- a/src/app/main.py
+++ b/src/app/main.py
@@ -1,14 +1,44 @@
 import gradio as gr
 import requests
 
-def echo(message, history):
+is_gpt3 = False
+
+
+def echo_gptchatbot(message, history):
+    # Post the message to the server
+    response = requests.post(
+        "http://127.0.0.1:8000/api/chat/gpt3", data={"user_request": message}
+    )
+    # Return the response
+    llm_output = response.json()["result"]["content"]
+
+    return llm_output
+
+
+def echo_llamachatbot(message, history):
     # Post the message to the server
-    response = requests.post("http://127.0.0.1:8000/api/chat/gpt3", data={"user_request": message})
+    response = requests.post(
+        "http://127.0.0.1:8000/api/chat/llama", data={"user_request": message}
+    )
+
     # Return the response
-    llm_output = response.json()["result"]['content']
-    
+    llm_output = response.json()["result"][0]["generated_text"]
+
     return llm_output
-    
-demo = gr.ChatInterface(fn=echo, examples=["What is OpenAI?", "What is LLM?"], title="LLM Chatbot")
 
-demo.launch()
\ No newline at end of file
+
+# Create a Gradio interface with the chatbot
+if is_gpt3:
+    demo = gr.ChatInterface(
+        fn=echo_gptchatbot,
+        examples=["What is OpenAI?", "What is GPT-3?"],
+        title="GPT-3 Chatbot",
+    )
+else:
+    demo = gr.ChatInterface(
+        fn=echo_llamachatbot,
+        examples=["What is OpenAI?", "What is LLM?"],
+        title="LLM Chatbot - Llama 2 7B",
+    )
+
+demo.launch()
diff --git a/src/index/main.py b/src/index/main.py
index 444de64..e1dbf8a 100644
--- a/src/index/main.py
+++ b/src/index/main.py
@@ -9,8 +9,8 @@
 if isProduction:
     app = FastAPI(
         title="LLM API Endpoints",
-        docs_url=None,                  # Disable docs (Swagger UI)
-        redoc_url=None,                 # Disable redoc
+        docs_url=None,  # Disable docs (Swagger UI)
+        redoc_url=None,  # Disable redoc
     )
 else:
     app = FastAPI(title="LLM API Endpoints")
@@ -40,6 +40,7 @@ async def gpt_chat(user_request: str = Form(...)):
 
     return {"result": result}
 
+
 @app.post("/api/chat/llama", tags=["Llama 2 7B Chat"])
 async def llama_chat(user_request: str = Form(...)):
     """