Update app.py

app.py

@@ -1,9 +1,48 @@
+from langchain.prompts import ChatPromptTemplate
+from langchain_community.document_loaders import JSONLoader
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain_cohere import ChatCohere
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnableLambda, RunnablePassthrough
+
+
+embedding_function = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+
+loader = JSONLoader(file_path="games.json", jq_schema=".games[]", text_content=False)
+documents = loader.load()
+
+db = Chroma.from_documents(documents, embedding_function)
+
+retriever = db.as_retriever(
+    search_type="mmr",
+    search_kwargs={'k': 500, 'fetch_k': 500}
+)
+
+template = """Answer the question based only on the following context:
+{context}
+
+Question: {question}
+"""
+prompt = ChatPromptTemplate.from_template(template)
+
+COHERE_API_KEY = os.getenv("COHERE_API_KEY")
+
+model = ChatCohere()
+
+chain = (
+    {"context": retriever, "question": RunnablePassthrough()}
+    | prompt
+    | model
+    | StrOutputParser()
+)
+
+
 import gradio as gr
 import os
 import cohere
 
 
-COHERE_API_KEY = os.getenv("COHERE_API_KEY")
 client = cohere.ClientV2(COHERE_API_KEY)
 
 COHERE_MODEL = "command-r-plus"
@@ -16,22 +55,29 @@ def respond(
     top_p,
 ):
 
-    system_message = '''
-    You are a friendly video game recommendation expert chatbot. 
+    query = message
+    retrieved_response = chain.invoke(query)
+
+    system_message = f"""
+    You are a friendly video game recommendation expert chatbot.
     Your task is to help parents and guardians to find appropriate video games for their children.
     Extract the child's age, preferred genre and multiplayer preference.
 
-    After you extracted the information you need, you should: 
-    - Suggest 5 video games that fit the given criteria.  
-    - If no games exactly match the genre, suggest similar alternatives.  
+    After you extracted the information you need, you should:
+    - Suggest 5 video games that fit the given criteria.
+    - If no games exactly match the genre, suggest similar alternatives.
     - If multiplayer is required, only include games with co-op mode.
 
+    Use the following game info to generate your suggestions:
+    {retrieved_response}
+    If you don't find enough games in the info you are given, then use your own knowledge.
+
     ### Response Format:
-    Game 1:  
-    - Name: [Game Title]  
-    - Genre: [Genre]  
-    - Age Suitability: [Age Group]  
-    - Multiplayer: [Yes/No]  
+    Game 1:
+    - Name: [Game Title]
+    - Genre: [Genre]
+    - Age Suitability: [Age Group]
+    - Multiplayer: [Yes/No]
     - Description: [Short summary of the game]
     - The reasons why you recommend the game
 
@@ -45,19 +91,18 @@ def respond(
     If they are not satisfied, then give the user the options of receiving more recommendations or changing their preferences.
     If they have questions about a game/games then provide the user with real information about the game/games.
     If they are satisifed and have no questions, then tell them that you were very happy to help and end the conversation.
-    '''
+    """
     
-    messages = [{"role": "system", "content": system_message}]
-
-    messages.append({"role": "user", "content": message})
+    messages = [{"role": "system", "content": system_message},
+                  {"role": "user", "content": message}]
 
     response = ""
-    
+
     response = client.chat(
-            messages=messages,
-            model=COHERE_MODEL,
-            temperature=temperature,
-            max_tokens=max_tokens,
+          messages=messages,
+          model=COHERE_MODEL,
+          temperature=temperature,
+          max_tokens=max_tokens,
     )
     yield response.message.content[0].text