agent + gr blocks

.gitignore

@@ -1,3 +1,5 @@
 .env
 hf_to_chroma_ds
-__pycache__
+__pycache__
+app_archive.py
+test_app.ipynb

README.md

@@ -4,7 +4,7 @@ emoji: 📚
 colorFrom: red
 colorTo: purple
 sdk: gradio
-sdk_version: 4.44.0
+sdk_version: 5.5.0
 app_file: app.py
 pinned: false
 startup_duration_timeout: 1h

app.py

@@ -1,166 +1,81 @@
-import os
-from dotenv import load_dotenv
-
+import time
 import gradio as gr
-
-from langchain_chroma import Chroma
-from langchain.prompts import ChatPromptTemplate
-from langchain.chains import create_retrieval_chain, create_history_aware_retriever
-from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain_core.prompts import MessagesPlaceholder
-from langchain_community.chat_message_histories import ChatMessageHistory
-from langchain_core.runnables.history import RunnableWithMessageHistory
-from langchain_core.documents import Document
-from langchain_core.retrievers import BaseRetriever
-from langchain_core.callbacks import CallbackManagerForRetrieverRun
-from langchain_core.vectorstores import VectorStoreRetriever
-from langchain_openai import ChatOpenAI
-from langchain.callbacks.tracers import ConsoleCallbackHandler
-from langchain_huggingface import HuggingFaceEmbeddings 
-
-from datasets import load_dataset
-import chromadb
-from typing import List
-from mixedbread_ai.client import MixedbreadAI
-from tqdm import tqdm
-
-# Global params
-CHROMA_PATH = "chromadb_mem10_mxbai_800_complete"
-MODEL_EMB = "mxbai-embed-large"
-MODEL_RRK = "mixedbread-ai/mxbai-rerank-large-v1"
-LLM_NAME = "gpt-4o-mini"
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-MXBAI_API_KEY = os.environ.get("MXBAI_API_KEY")
-HF_TOKEN = os.environ.get("HF_TOKEN")
-HF_API_KEY = os.environ.get("HF_API_KEY")
-
-# MixedbreadAI Client
-# device = "cuda:0" if torch.cuda.is_available() else "cpu"
-mxbai_client = MixedbreadAI(api_key=MXBAI_API_KEY)
-model_emb = "mixedbread-ai/mxbai-embed-large-v1"  
-
-# Set up ChromaDB
-memoires_ds = load_dataset("eliot-hub/memoires_vec_800", split="data", token=HF_TOKEN, streaming=True)
-batched_ds = memoires_ds.batch(batch_size=41000)
-client = chromadb.Client()
-collection = client.get_or_create_collection(name="embeddings_mxbai") 
-
-for batch in tqdm(batched_ds, desc="Processing dataset batches"): 
-    collection.add(
-        ids=batch["id"],
-        metadatas=batch["metadata"],
-        documents=batch["document"],
-        embeddings=batch["embedding"],      
+from tools import create_agent
+from langchain_core.messages import RemoveMessage
+from langchain_core.messages import trim_messages
+# from toolkits import create_agent
+# from langchain.schema import AIMessage, HumanMessage, SystemMessage
+
+AGENT = create_agent()
+theme = gr.themes.Default(primary_hue="red", secondary_hue="red")
+
+def filter_msg(msg_list:list, keep_n:int) -> list:
+    """Keep only last keep_n messages from chat history. Preserves structure user msg -> tool msg -> ai msg"""
+    msg = trim_messages(
+            msg_list,
+            strategy="last",
+            token_counter=len,
+            max_tokens=keep_n,
+            start_on="human",
+            end_on=("tool", "ai"),
+            include_system=True,
     )
-print(f"Collection complete: {collection.count()}")
-
-db = Chroma(
-    client=client,
-    collection_name=f"embeddings_mxbai",
-    embedding_function = HuggingFaceEmbeddings(model_name=model_emb) 
-)
-
-
-# Reranker class
-class Reranker(BaseRetriever):
-    retriever: VectorStoreRetriever
-    # model: CrossEncoder
-    k: int
-
-    def _get_relevant_documents(
-        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
-    ) -> List[Document]:
-        docs = self.retriever.invoke(query)
-        results = mxbai_client.reranking(model=MODEL_RRK, query=query, input=[doc.page_content for doc in docs], return_input=True, top_k=self.k)
-        return [Document(page_content=res.input) for res in results.data]
-
-# Set up reranker + LLM
-retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 25})
-reranker = Reranker(retriever=retriever, k=4)  #Reranker(retriever=retriever, model=model, k=4)
-llm = ChatOpenAI(model=LLM_NAME, verbose=True) #, api_key=OPENAI_API_KEY, )
-
-# Set up the contextualize question prompt
-contextualize_q_system_prompt = (
-    "Compte tenu de l'historique des discussions et de la dernière question de l'utilisateur "
-    "qui peut faire référence à un contexte dans l'historique du chat, "
-    "formuler une question autonome qui peut être comprise "
-    "sans l'historique du chat. Ne répondez PAS à la question, "
-    "juste la reformuler si nécessaire et sinon la renvoyer telle quelle."
-)
-
-contextualize_q_prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", contextualize_q_system_prompt),
-        MessagesPlaceholder("chat_history"),
-        ("human", "{input}"),
-    ]
-)
-
-# Create the history-aware retriever
-history_aware_retriever = create_history_aware_retriever(
-    llm, reranker, contextualize_q_prompt
-)
-
-# Set up the QA prompt
-system_prompt = (
-    "Réponds à la question en te basant uniquement sur le contexte suivant: \n\n {context}"
-    "Si tu ne connais pas la réponse, dis que tu ne sais pas."
-)
-qa_prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", system_prompt),
-        MessagesPlaceholder("chat_history"),
-        ("human", "{input}"),
-    ]
-)
-
-# Create the question-answer chain
-question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
-rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
-
-# Set up the conversation history
-store = {}
-
-def get_session_history(session_id: str) -> ChatMessageHistory:
-    if session_id not in store:
-        store[session_id] = ChatMessageHistory()
-    return store[session_id]
-
-conversational_rag_chain = RunnableWithMessageHistory(
-    rag_chain,
-    get_session_history,
-    input_messages_key="input",
-    history_messages_key="chat_history",
-    output_messages_key="answer",
-)
-
-# Gradio interface
-def chatbot(message, history):
-    session_id = "gradio_session"
-    response = conversational_rag_chain.invoke(
-        {"input": message},
-        config={
-            "configurable": {"session_id": session_id},
-            "callbacks": [ConsoleCallbackHandler()]
-        },
-    )["answer"]
-    return response
-
-iface = gr.ChatInterface(
-    chatbot,
-    title="Dataltist Chatbot",
-    description="Posez vos questions sur l'assurance",
-    textbox=gr.Textbox(placeholder="Qu'est-ce que l'assurance multirisque habitation ?", container=False, scale=9),
-    theme="soft",
-    # examples=[
-    #     "Qu'est-ce que l'assurance multirisque habitation ?",
-    #     "Qu'est-ce que la garantie DTA ?",
-    # ],
-    retry_btn=None,
-    undo_btn=None,
-    submit_btn=gr.Button(value="Envoyer", icon="./send_icon.png", variant="primary"),
-    clear_btn="Effacer la conversation",
-)
+    return [m.id for m in msg]
+
+def agent_response(query, config, keep_n=10):
+    messages = AGENT.get_state(config).values.get("messages", [])
+    
+    if len(messages) > keep_n:
+        keep_msg_ids = filter_msg(messages, keep_n)
+        AGENT.update_state(config, {"messages": [RemoveMessage(id=m.id) for m in messages if m.id not in keep_msg_ids]})
+        print("msg removed")
+
+    # Generate answer
+    answer = AGENT.invoke({"messages":query}, config=config)
+    return answer["messages"][-1].content
+
+
+js_func = """
+function refresh() {
+    const url = new URL(window.location);
+
+    if (url.searchParams.get('__theme') != 'light') {
+        url.searchParams.set('__theme', 'light');
+        window.location.href = url.href;
+    }
+}
+"""
+
+
+def delete_agent():
+    print("del agent")
+    global AGENT
+    AGENT = create_agent()
+    # print(AGENT.get_state(config).values.get("messages"), "\n\n")
+
+with gr.Blocks(theme=theme, js=js_func, title="Dataltist", fill_height=True) as iface:
+    gr.Markdown("# Dataltist Chatbot 🚀")
+    chatbot = gr.Chatbot(show_copy_button=True, show_share_button=False, type="messages", scale=1)
+    msg = gr.Textbox(lines=1, show_label=False, placeholder="Posez vos questions sur l'assurance")  #  submit_btn=True
+    clear = gr.ClearButton([msg, chatbot], value="Effacer 🗑")
+    config = {"configurable": {"thread_id": "1"}} 
+
+
+    def user(user_message, history: list):
+        return "", history + [{"role": "user", "content": user_message}]
+    
+    def bot(history: list):
+        bot_message = agent_response(history[-1]["content"], config)  #AGENT.invoke({"messages":history[-1]["content"]}, config=config)
+        history.append({"role": "assistant", "content": ""})
+        for character in bot_message:
+            history[-1]['content'] += character
+            # time.sleep(0.005)
+            yield history
+
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, chatbot, chatbot
+    )
+    iface.unload(delete_agent)
 
 if __name__ == "__main__":
-    iface.launch()  # share=True
+    iface.launch()  # share=True # auth=("admin", "admin")

tools.py

@@ -0,0 +1,126 @@
+from langchain_community.tools import TavilySearchResults
+
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
+from langchain_core.vectorstores import VectorStoreRetriever
+from langgraph.prebuilt import create_react_agent
+from langchain_core.documents import Document
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from mixedbread_ai.client import MixedbreadAI
+from langchain.chains import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.prompts import ChatPromptTemplate
+from dotenv import load_dotenv
+import os
+from langchain_chroma import Chroma
+import chromadb
+from typing import List
+from datasets import load_dataset
+from langchain_huggingface import HuggingFaceEmbeddings 
+
+
+load_dotenv()
+# Global params
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+MODEL_EMB = "mxbai-embed-large"
+MODEL_RRK = "mixedbread-ai/mxbai-rerank-large-v1"
+LLM_NAME = "gpt-4o-mini"
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+MXBAI_API_KEY = os.environ.get("MXBAI_API_KEY")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_API_KEY = os.environ.get("HF_API_KEY")
+
+# MixedbreadAI Client
+mxbai_client = MixedbreadAI(api_key=MXBAI_API_KEY)
+model_emb = "mixedbread-ai/mxbai-embed-large-v1"  
+
+# # Set up ChromaDB
+memoires_ds = load_dataset("eliot-hub/memoires_vec_800", split="data", token=HF_TOKEN, streaming=True)
+batched_ds = memoires_ds.batch(batch_size=41000)
+client = chromadb.Client()
+collection = client.get_or_create_collection(name="embeddings_mxbai") 
+
+
+
+llm_4o = ChatOpenAI(model="gpt-4o-mini", api_key=OPENAI_API_KEY, temperature=0)
+
+    
+
+def init_rag_tool():
+    """Init tools to allow an LLM to query the documents"""
+    # client = chromadb.PersistentClient(path=CHROMA_PATH)
+    db = Chroma(
+        client=client,
+        collection_name=f"embeddings_mxbai",
+        embedding_function = HuggingFaceEmbeddings(model_name=model_emb) 
+    )
+    # Reranker class
+    class Reranker(BaseRetriever):
+        retriever: VectorStoreRetriever
+        # model: CrossEncoder
+        k: int
+
+        def _get_relevant_documents(
+            self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+        ) -> List[Document]:
+            docs = self.retriever.invoke(query)
+            results = mxbai_client.reranking(model=MODEL_RRK, query=query, input=[doc.page_content for doc in docs], return_input=True, top_k=self.k)
+            return [Document(page_content=res.input) for res in results.data]
+
+    # Set up reranker + LLM
+    retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 25})
+    reranker = Reranker(retriever=retriever, k=4)  #Reranker(retriever=retriever, model=model, k=4)
+    llm = ChatOpenAI(model=LLM_NAME, verbose=True) 
+
+
+    system_prompt = (
+    "Réponds à la question en te basant uniquement sur le contexte suivant: \n\n {context}"
+    "Si tu ne connais pas la réponse, dis que tu ne sais pas."
+    )
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            ("human", "{input}"),
+        ]
+    )
+
+    question_answer_chain = create_stuff_documents_chain(llm, prompt)
+    rag_chain = create_retrieval_chain(reranker, question_answer_chain)
+
+    rag_tool = rag_chain.as_tool(
+        name="RAG_search",
+        description="Recherche d'information dans les mémoires d'actuariat",
+        arg_types={"input": str},
+    )
+    return rag_tool
+
+
+def init_websearch_tool():
+    web_search_tool = TavilySearchResults(
+        name="Web_search",
+        max_results=5,
+        description="Recherche d'informations sur le web",
+        search_depth="advanced",
+        include_answer=True,
+        include_raw_content=True,
+        include_images=False,
+        verbose=False,
+    )
+    return web_search_tool
+
+
+def create_agent():
+    rag_tool = init_rag_tool()
+    web_search_tool = init_websearch_tool()
+    memory = MemorySaver()  
+    llm_4o = ChatOpenAI(model="gpt-4o-mini", api_key=OPENAI_API_KEY, verbose=True, temperature=0, streaming=True)
+    tools = [rag_tool, web_search_tool]
+    system_message = """
+        Tu es un assistant dont la fonction est de répondre à des questions à propos de l'assurance et de l'actuariat.
+        Utilise les outils RAG_search ou Web_search pour répondre aux questions de l'utilisateur.
+    """    # Dans la réponse finale, sépare les informations de l'outil RAG et de l'outil Web. 
+
+    react_agent = create_react_agent(llm_4o, tools, state_modifier=system_message, checkpointer=memory, debug=False)
+    return react_agent