Update app.py

app.py

@@ -1,70 +1,245 @@
+import glob
 import gradio as gr
-from huggingface_hub import InferenceClient
+from langchain_community.document_loaders.csv_loader import CSVLoader
+from langchain_community.document_loaders import Docx2txtLoader, TextLoader, PyPDFLoader
+from langchain_text_splitters import CharacterTextSplitter, RecursiveCharacterTextSplitter, TokenTextSplitter
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from huggingface_hub import snapshot_download, upload_folder
 
+from langchain.tools import tool
+from langchain.agents import create_agent
+from langchain.agents.middleware import dynamic_prompt, ModelRequest
 
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+snapshot_download(repo_id="CGIAR/weai-ref", 
+                  repo_type="dataset", 
+                  token=os.getenv('HF_TOKEN'), 
+                  local_dir='./refs'
+                  )
 
-    messages = [{"role": "system", "content": system_message}]
+snapshot_download(repo_id="CGIAR/weai-docs", 
+                  repo_type="dataset", 
+                  token=os.getenv('HF_TOKEN'), 
+                  local_dir='./docs'
+                  )
+warnings.filterwarnings('ignore')
+os.environ["WANDB_DISABLED"] = "true"
 
-    messages.extend(history)
+repo_id = "meta-llama/Llama-3.3-70B-Instruct"
 
-    messages.append({"role": "user", "content": message})
+model = HuggingFaceEndpoint(
+    task='conversational',
+    repo_id = repo_id,
+    temperature = 0.5,
+    huggingfacehub_api_token=os.getenv('HF_TOKEN'),
+    max_new_tokens = 1500,
+)
+
+chat_llm = ChatHuggingFace(llm=model, verbose=True)
+
+model_name = "sentence-transformers/all-mpnet-base-v2"
+model_kwargs = {"device": "cuda"}
+
+embeddings = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs)
+
+def docs_return(directory_path, flag):
+    docx_file_pattern = '*.docx'
+    pdf_file_pattern = '*.pdf'
+    txt_file_pattern = '*.txt'
+
+    docx_file_paths = glob.glob(directory_path + docx_file_pattern)
+    pdf_file_paths = glob.glob(directory_path + pdf_file_pattern)
+    txt_file_paths = glob.glob(directory_path + txt_file_pattern)
+
+    all_doc, all_doc2 = [], []
+
+    for x in docx_file_paths:
+        loader = Docx2txtLoader(x)
+        documents = loader.load()
+        all_doc.extend(documents)
+        all_doc2.append(str(documents[0].page_content))
+
+    for x in pdf_file_paths:
+        loader = PyPDFLoader(x, extract_images=True)
+        docs_lazy = loader.lazy_load()
+        documents = []
+        for doc in docs_lazy:
+            documents.append(doc)
+        all_doc.extend(documents)
+        all_doc2.append(str(documents[0].page_content))
+
+    for x in txt_file_paths:
+        loader = TextLoader(x)
+        documents = loader.load()
+        all_doc.extend(documents)
+        all_doc2.append(str(documents[0].page_content))
+
+    docs = '\n\n'.join(all_doc2)
+
+    return all_doc if flag == 0 else docs
+
+def get_text_splitter(splitter_type='character', 
+                      chunk_size=500, 
+                      chunk_overlap=30, 
+                      separator="\n", 
+                      max_tokens=1000):
+    if splitter_type == 'character':
+        return CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap, separator=separator)
+    elif splitter_type == 'recursive':
+        return RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    elif splitter_type == 'token':
+        return TokenTextSplitter(chunk_size=max_tokens, chunk_overlap=chunk_overlap)
+    else:
+        raise ValueError("Unsupported splitter type. Choose from 'character', 'recursive', or 'token'.")
+
+splitter_type='character' 
+chunk_size=1500
+chunk_overlap=30
+separator="\n"
+max_tokens=1000
+docs_path = "./docs/"
+
+all_doc = docs_return(docs_path, 0)
+
+# Use the splitter parameters
+text_splitter = get_text_splitter(splitter_type=splitter_type, 
+                                  chunk_size=chunk_size, 
+                                  chunk_overlap=chunk_overlap, 
+                                  separator=separator, 
+                                  max_tokens=max_tokens)
+
+# Split the documents using the text splitter
+docs = text_splitter.split_documents(documents=all_doc)
+
+# Create a Chroma vector database
+docs_vector_db = Chroma.from_documents(docs, embeddings, persist_directory="chroma_data")
+
+REFS_CSV_PATH = f"{DATA_DIR}/WEAI reference list - Sheet1.csv"
+REFS_CHROMA_PATH = "./refs/"
+
+loader = CSVLoader(file_path=REFS_CSV_PATH,
+                   source_column="Description (what it contains and what it's useful for)")
+refs = loader.load()
 
-    response = ""
+refs_vector_db = Chroma.from_documents(
+    refs, embeddings, persist_directory=REFS_CHROMA_PATH
+)
+
+@dynamic_prompt
+def ref_context(request: ModelRequest) -> str:
+    """Inject context into state messages."""
+    last_query = request.state["messages"][-1].text
+    ref_content = refs_vector_db.as_retriever(k=10)
+
+    system_message = (
+        """Your job is to use relevant links and email addresses to 
+        direct users to in order to reach and contact the WEAI team. If you don't know 
+        an answer, say you don't know. Do not state that you are referring to the 
+        provided context and respond as if you were in charge of the WEAI helpdesk.""" 
+        f"\n\n{ref_content}"
+    )
+
+    return system_message
+
+contact_agent = (create_agent(chat, tools=[], middleware=[ref_context]))
+
+@tool("contact", description="refer users to WEAI team using links and contact details")
+def call_contact_agent(query: str):
+    result = contact_agent.invoke({"messages": [{"role": "user", "content": query}]})
+    return result["messages"][-1].content
+    
+
+@dynamic_prompt
+def doc_context(request: ModelRequest) -> str:
+    """Inject context into state messages."""
+    last_query = request.state["messages"][-1].text
+    doc_content = docs_vector_db.as_retriever(k=10)
+
+    system_message = (
+        """Your job is to use resources from the International Food
+            Policy Research Institute to answer questions about women empowerment in agriculture.
+            Use the following context to answer questions. Be as detailed
+            as possible, but don't make up any information that's not
+            from the context and where possible reference related studies and resources as examples
+            from the context you have. If you don't know an answer, say you don't know. 
+            Be concise but thorough in your response and try not to exceed the output token limit.
+            Do not state that you are referring to the provided context and respond 
+            as if you were in charge of the WEAI helpdesk. """ 
+        f"\n\n{doc_content}"
+    )
+
+    return system_message
+
+support_agent = (create_agent(chat, tools=[call_contact_agent], middleware=[doc_context]))
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
 
-        response += token
-        yield response
+@tool("support", description="respond to user queries using context in WEAI docs")
+def call_support_agent(query: str):
+    result = support_agent.invoke({"messages": [{"role": "user", "content": query}]})
+    return result["messages"][-1].content
+    
+support_instructions = """
+You are in charge of the WEAI helpdesk.
+Your job is to answer user queries using provided context and references 
+and refer users to WEAI personnel as well as relevant resource links where necessary.
 
+Steps:
+1. Use the support tool to answer queries to the best of your knowledge.
+2. If no contact information or links are provided in the response, use the 
+   contact tool to add all relevant contact and resource information to the response.
+3. Return only a complete response with included contact and resource information.
+"""
 
+response_agent = create_agent(model=chat, 
+                              tools=[call_contact_agent, call_support_agent],
+                              system_prompt=support_instructions,
+                              )
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
 with gr.Blocks() as demo:
     with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+        gr.LoginButton()   
+    gr.Markdown("# WEAI-bot")
+    chatbot = gr.Chatbot(type='messages', 
+                         allow_tags=True)
+    msg = gr.Textbox()
+    clear = gr.ClearButton([msg, chatbot])
+
+    def handle_undo(history, undo_data: gr.UndoData):
+        return history[:undo_data.index], history[undo_data.index]['content'][0]["text"]
+
+    def handle_retry(history, retry_data: gr.RetryData):
+        new_history = history[:retry_data.index]
+        previous_prompt = history[retry_data.index]['content'][0]["text"]
+        yield from respond(previous_prompt, new_history)
+
+    def support_agent_fn(message, history):
+        result = support_agent.invoke({"messages": [{"role": "user", "content": message}]})
+
+        response = result['messages'][-1].content#.split('<|start_header_id|>assistant<|end_header_id|>')[-1].strip()
+        history.append({"role": "user", "content": message})
+        history.append({"role": "assistant", "content": response})
+
+        return response, history
+        
+    def handle_like(data: gr.LikeData):
+        if data.liked:
+            print("You upvoted this response: ", data.value)
+        else:
+            print("You downvoted this response: ", data.value)
+
+    def handle_edit(history, edit_data: gr.EditData):
+        new_history = history[:edit_data.index]
+        new_history[-1]['content'] = [{"text": edit_data.value, "type": "text"}]
+        return new_history
+
+    msg.submit(support_agent_fn, [msg, chatbot], [msg, chatbot])
 
+    chatbot.undo(handle_undo, chatbot, [chatbot, msg])
+    chatbot.retry(handle_retry, chatbot, chatbot)
+    chatbot.like(handle_like, None, None)
+    chatbot.edit(handle_edit, chatbot, chatbot)
 
 if __name__ == "__main__":
     demo.launch()