Update app.py

app.py

@@ -1,10 +1,5 @@
-import gradio as gr
-from mistralai.client import MistralClient
-from mistralai.models.chat_completion import ChatMessage
-import requests
-import numpy as np
-import faiss
 import os
+import gradio as gr
 from llama_index import VectorStoreIndex, SimpleDirectoryReader, SummaryIndex
 from llama_index.readers import SimpleWebPageReader
 from llama_index.llms import MistralAI
@@ -12,89 +7,53 @@ from llama_index.embeddings import MistralAIEmbedding
 from llama_index import ServiceContext
 from llama_index.query_engine import RetrieverQueryEngine
 
-title = "Gaia Mistral Chat Demo with RAG"
+
+title = "Gaia Mistral Chat Demo"
 description = "Exemple d'assistant avec Gradio et Mistral AI via son API"
 placeholder = "Posez moi une question sur l'agriculture"
 placeholder_url = "Donner moi une url qui va servir de contexte agricole complémentaire"
+placeholder_api_key = "API key"
 examples = ["Comment fait on pour produire du maïs ?", "Rédige moi une lettre pour faire un stage dans une exploitation agricole", "Comment reprendre une exploitation agricole ?"]
 
-api_key = os.environ.get("MISTRAL_API_KEY")
-client = MistralClient(api_key=api_key)
-chat_model = 'mistral-small'
+query_engine = None
 
-llm = MistralAI(api_key=api_key,model="mistral-medium")
-embed_model = MistralAIEmbedding(model_name='mistral-embed', api_key=api_key)
-service_context = ServiceContext.from_defaults(chunk_size=512, llm=llm, embed_model=embed_model)
+with gr.Blocks() as demo:
 
-# build a vector database with documents
-def setup_db_with_url(url):
-    global query_engine
-    documents = SimpleWebPageReader(html_to_text=True).load_data([url])
-    # insert in DB
-    index = VectorStoreIndex.from_documents(documents, service_context=service_context)
-    query_engine = index.as_query_engine(similarity_top_k=2)
-    # get document source
-    # response = requests.get(url)
-    # text = response.text
-    # print(text)
-    # chunk_size = 512
-    # split in chunks
-    # chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
-    # embbed in Mistral to have vectors
-    # text_embeddings = np.array([get_text_embedding(chunk) for chunk in chunks])
-    # print(text_embeddings)
-    # d = text_embeddings.shape[1]
-    documents
-    
+    gr.Markdown(""" ### Welcome to Level 2 Demo 
+    Add an url and your API key at the bottom of the bottom before interacting with the Chat. This demo allows you to interact with a webpage and then ask questions to Mistral APIs. Mistral will answer with the context extracted from the webpage.   
+    """)
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    clear = gr.ClearButton([msg, chatbot])
 
-# vector data with Mistral
-def get_text_embedding(input):
-    embeddings_batch_response = client.embeddings(
-          model="mistral-embed",
-          input=input
-      )
-    return embeddings_batch_response.data[0].embedding
-
-# build a prompt
-def build_prompt(user_input):
-    retrieved_chunk = query_engine.query(user_input)
-    prompt = f"""
-        Context information is below.
-        ---------------------
-        {retrieved_chunk}
-        ---------------------
-        Given the context information and not prior knowledge, answer the query.
-        Query: {user_input}
-        Answer:
-        """
-
-def chat_with_mistral(user_input, history):
-    prompt = build_prompt(user_input)
-    messages = [ChatMessage(role="user", content=prompt)]
+    with gr.Row():
+        api_key_text_box = gr.Textbox(placeholder=placeholder_api_key, container=False, scale=7)
 
-    chat_response = client.chat(model=chat_model, messages=messages)
-    mistral_content = chat_response.choices[0].message.content
-    histories = history + [(mistral_content, None)]
-    return [mistral_content, histories]
+    def setup_with_url(url, api_key):
+        global query_engine
+        # Set-up clients
+        llm = MistralAI(api_key=api_key,model="mistral-medium")
+        embed_model = MistralAIEmbedding(model_name='mistral-embed', api_key=api_key)
+        service_context = ServiceContext.from_defaults(chunk_size=1024, llm=llm, embed_model=embed_model)
 
-with gr.Blocks() as iface:
+        # Set-up db
+        documents = SimpleWebPageReader(html_to_text=True).load_data([url])
+        index = VectorStoreIndex.from_documents(documents, service_context=service_context)
+        query_engine = index.as_query_engine(similarity_top_k=15)
+        return "I'm ready, please add a question here."
 
-    with gr.Row():
-        gr.Markdown("#Mixtral-8x7B Playground Space!")
     with gr.Row():
         url_msg = gr.Textbox(placeholder=placeholder_url, container=False, scale=7)
-        url_btn = gr.Button(value="🔄", interactive=True)
-    with gr.Row():
-        url_return = gr.Textbox(value='', container=False, scale=7)
-        url_btn.click(setup_db_with_url, url_msg, url_return)
-    with gr.Row():
-        chatbot=gr.Chatbot(height=300)
-    with gr.Row():
-        msg = gr.Textbox(placeholder=placeholder, container=False, scale=7)
-        msg_btn = gr.Button("Envoyer")
+        url_btn = gr.Button(value="Set-up API and process url ✅", interactive=True)
+        url_btn.click(setup_with_url, [url_msg, api_key_text_box], msg, show_progress= "full")
+
+    
 
-    msg_btn.click(chat_with_mistral, [msg, chatbot], [msg, chatbot] )
+    def respond(message, chat_history):
+        response = query_engine.query(message)
+        chat_history.append((message, str(response)))
+        return chat_history
 
-iface.title = title
+    msg.submit(respond, [msg, chatbot], [chatbot])
 
-iface.launch(share=True)
+demo.launch()