[Release] Docs Agent version 0.2.2

What's changed:

- Bug fixes related to the new AQA model integration.
  - Prevent Docs Agent from crashing when OAuth is not set
    and the AQA model is not used.
  - Prevent Docs Agent from crashing when the AQA model fails to respond.
- Enable the Docs Agent UI to display the URL matadata from an entry
  used by the AQA model.
- Log the `Answerable probability` score when the AQA model is used.
- Update `whats-new.md` for January, 2024.

Author: kyolee415

demos/palm/python/docs-agent/README.md

@@ -374,6 +374,38 @@ The following issues have been identified and need to be worked on:
 
 This section provides instructions on how to set up the Docs Agent project on a Linux host machine.
 
+### 0. (Optional) Authorize credentials for Docs Agent
+
+**This step is needed only if you plan to use Gemini's AQA model.** For more information on this
+feature, see the
+[Using the Semantic Retrieval API and AQA model](#using-the-semantic-retrieval-api-and-aqa-model)
+section above.
+
+1. Download the `client_secret.json` file from your Google Cloud Project (GCP) account.
+
+   See [Authorize credentials for a desktop application][authorize-credentials]
+   on the _AI for Developers_ doc site.
+
+2. Copy the `client_secret.json` file to your host machine.
+
+3. To authenticate credentials, run the following command in the directory of
+   the host machine where the `client_secret.json` file is located:
+
+   ```
+   gcloud auth application-default login --client-id-file=client_secret.json --scopes='https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/generative-language.retriever'
+   ```
+
+   This command opens a browser and asks to log in using your Google account.
+
+   **Note**: If the `gcloud` command doesn’t exist, install the Google Cloud SDK
+   on your host machine: `sudo apt install google-cloud-sdk`
+
+4. Follow the instructions on the browser and click **Allow** to authenticate.
+
+   This saves the authenticated credentials for Docs Agent
+   (`application_default_credentials.json`) in the `$HOME/.config/gcloud/`
+   directory of your host machine.
+
 ### 1. Prerequisites
 
 1. Update the Linux package repositories on the host machine:
@@ -700,6 +732,7 @@ Meggin Kearney (`@Meggin`), and Kyo Lee (`@kyolee415`).
 [contribute-to-docs-agent]: #contribute-to-docs-agent
 [set-up-docs-agent]: #set-up-docs-agent
 [markdown-to-plain-text]: ./scripts/markdown_to_plain_text.py
+[files-to-plain-text]: ./scripts/files_to_plain_text.py
 [populate-vector-database]: ./scripts/populate_vector_database.py
 [context-source-01]: http://eventhorizontelescope.org
 [fact-check-section]: #using-a-language-model-to-fact-check-its-own-response
@@ -722,3 +755,4 @@ Meggin Kearney (`@Meggin`), and Kyo Lee (`@kyolee415`).
 [aqa-model]: https://ai.google.dev/models/gemini#model_variations
 [oauth-quickstart]: https://ai.google.dev/docs/oauth_quickstart
 [inline-passages]: https://ai.google.dev/docs/semantic_retriever#more_options_aqa_using_inline_passages
+[authorize-credentials]: https://ai.google.dev/docs/oauth_quickstart#authorize-credentials

demos/palm/python/docs-agent/aqa.py

@@ -86,10 +86,15 @@ def create_a_doc(self, corpus_name, page_title, page_url):
             document_resource_name = get_document_response.name
         return document_resource_name
 
-    def create_a_chunk(self, doc_name, text):
+    def create_a_chunk(self, doc_name, text, url):
         response = ""
         try:
+            # Create a chunk.
             chunk = glm.Chunk(data={"string_value": text})
+            # Add metadata.
+            chunk.custom_metadata.append(
+                glm.CustomMetadata(key="url", string_value=url)
+            )
             create_chunk_requests = []
             create_chunk_requests.append(
                 glm.CreateChunkRequest(parent=doc_name, chunk=chunk)
@@ -118,10 +123,10 @@ def create_a_chunk(self, doc_name, text):
                 else:
                     text_02 += line + "\n"
                 i += 1
-            self.create_a_chunk(doc_name, text_01)
-            self.create_a_chunk(doc_name, text_02)
+            self.create_a_chunk(doc_name, text_01, url)
+            self.create_a_chunk(doc_name, text_02, url)
         return response
 
     def create_a_doc_chunk(self, corpus_name, page_title, page_url, text):
         doc_name = self.create_a_doc(corpus_name, page_title, page_url)
-        return self.create_a_chunk(doc_name, text)
+        return self.create_a_chunk(doc_name, text, page_url)

demos/palm/python/docs-agent/chatbot/chatui.py

@@ -192,9 +192,16 @@ def ask_model(question):
     clickable_urls = markdown.markdown(
         query_result.fetch_formatted(Format.CLICKABLE_URL)
     )
-    fact_check_url = markdown.markdown(
-        query_result.fetch_nearest_formatted(Format.CLICKABLE_URL)
-    )
+    fact_check_url = ""
+    if docs_agent.check_if_aqa_is_used() and docs_agent.get_db_type() == "ONLINE_STORAGE":
+        aqa_response_url = docs_agent.get_aqa_response_url()
+        fact_check_url = (
+            '<a href="' + aqa_response_url + '">' + aqa_response_url + "</a>"
+        )
+    else:
+        fact_check_url = markdown.markdown(
+            query_result.fetch_nearest_formatted(Format.CLICKABLE_URL)
+        )
 
     ### PREPARE OTHER ELEMENTS NEEDED BY UI.
     # - Create a uuid for this request.
@@ -242,6 +249,21 @@ def ask_model(question):
 def parse_related_questions_response_to_html_list(response):
     soup = BeautifulSoup(response, "html.parser")
     for item in soup.find_all("li"):
+        # In case there are code tags, remove the tag and just replace with
+        # plain text
+        if item.find("code"):
+            text = item.find("code").text
+            item.code.replace_with(text)
+        # In case there are <p> tags within the <li> strip <p>
+        if item.find("p"):
+            text = item.find("p").text
+            link = soup.new_tag(
+                "a",
+                href=url_for("chatui.question", ask=urllib.parse.quote_plus(text)),
+            )
+            link.string = text
+            item.string = ""
+            item.append(link)
         if item.string is not None:
             link = soup.new_tag(
                 "a",
@@ -264,10 +286,19 @@ def log_question(uid, user_question, response):
     print("Question: " + user_question.strip() + "\n")
     print("Response:")
     print(response.strip() + "\n")
+    if docs_agent.check_if_aqa_is_used():
+        aqa_response = docs_agent.get_saved_aqa_response_json()
+        try:
+            probability = aqa_response.answerable_probability
+        except:
+            probability = 0.0
+        print("Answerable probability: " + str(probability) + "\n")
     with open("chatui_logs.txt", "a", encoding="utf-8") as log_file:
         log_file.write("[" + date.strftime(date_format) + "][UID " + str(uid) + "]\n")
         log_file.write("# " + user_question.strip() + "\n\n")
         log_file.write(response.strip() + "\n\n")
+        if docs_agent.check_if_aqa_is_used():
+            log_file.write("Answerable probability: " + str(probability) + "\n\n")
         log_file.close()
 
 

demos/palm/python/docs-agent/docs/whats-new.md

@@ -1,5 +1,9 @@
 # What's new in Docs Agent
 
+## January 2024
+
+* **Milestone: Docs Agent uses AQA model and Semantric Retrieval API**
+
 ## December 2023
 
 * **Milestone: Docs Agent uses Gemini model.**

demos/palm/python/docs-agent/docs_agent.py

@@ -93,9 +93,12 @@
 else:
     palm = PaLM(api_key=API_KEY, api_endpoint=PALM_API_ENDPOINT)
 
-embedding_function_gemini_retrieval = embedding_functions.GoogleGenerativeAiEmbeddingFunction(
-    api_key=API_KEY, model_name="models/embedding-001",
-    task_type="RETRIEVAL_QUERY")
+embedding_function_gemini_retrieval = (
+    embedding_functions.GoogleGenerativeAiEmbeddingFunction(
+        api_key=API_KEY, model_name="models/embedding-001", task_type="RETRIEVAL_QUERY"
+    )
+)
+
 
 class DocsAgent:
     """DocsAgent class"""
@@ -107,8 +110,9 @@ def __init__(self):
         )
         self.chroma = Chroma(LOCAL_VECTOR_DB_DIR)
         self.collection = self.chroma.get_collection(
-            COLLECTION_NAME, embedding_model=EMBEDDING_MODEL,
-            embedding_function=embedding_function_gemini_retrieval
+            COLLECTION_NAME,
+            embedding_model=EMBEDDING_MODEL,
+            embedding_function=embedding_function_gemini_retrieval,
         )
         # Update PaLM's custom prompt strings
         self.prompt_condition = CONDITION_TEXT
@@ -121,12 +125,21 @@ def __init__(self):
         self.is_aqa_used = IS_AQA_USED
         self.db_type = DB_TYPE
         # AQA model setup
-        self.generative_service_client = glm.GenerativeServiceClient()
-        self.retriever_service_client = glm.RetrieverServiceClient()
-        self.permission_service_client = glm.PermissionServiceClient()
+        self.generative_service_client = {}
+        self.retriever_service_client = {}
+        self.permission_service_client = {}
         self.corpus_display = PRODUCT_NAME + " documentation"
         self.corpus_name = "corpora/" + PRODUCT_NAME.lower().replace(" ", "-")
         self.aqa_response_buffer = ""
+        self.set_up_aqa_model_environment()
+
+    # Set up the AQA model environment
+    def set_up_aqa_model_environment(self):
+        if IS_AQA_USED == "YES":
+            self.generative_service_client = glm.GenerativeServiceClient()
+            self.retriever_service_client = glm.RetrieverServiceClient()
+            self.permission_service_client = glm.PermissionServiceClient()
+        return
 
     # Use this method for talking to a PaLM text model
     def ask_text_model_with_context(self, context, question):
@@ -203,7 +216,11 @@ def ask_aqa_model_using_local_vector_store(self, question):
         elif LOG_LEVEL == "DEBUG":
             self.print_the_prompt(verbose_prompt)
             print(aqa_response)
-        return aqa_response.answer.content.parts[0].text
+        try:
+            return aqa_response.answer.content.parts[0].text
+        except:
+            self.aqa_response_buffer = ""
+            return self.model_error_message
 
     # Use this method for talking to Gemini's AQA model using a corpus
     def ask_aqa_model_using_corpora(self, question):
@@ -243,7 +260,11 @@ def ask_aqa_model_using_corpora(self, question):
             self.print_the_prompt(verbose_prompt)
         elif LOG_LEVEL == "DEBUG":
             print(aqa_response)
-        return aqa_response.answer.content.parts[0].text
+        try:
+            return aqa_response.answer.content.parts[0].text
+        except:
+            self.aqa_response_buffer = ""
+            return self.model_error_message
 
     def ask_aqa_model(self, question):
         response = ""
@@ -323,6 +344,27 @@ def check_if_aqa_is_used(self):
     def get_saved_aqa_response_json(self):
         return self.aqa_response_buffer
 
+    # Retrieve the URL metadata from the AQA model's response
+    def get_aqa_response_url(self):
+        url = ""
+        try:
+            # Get the metadata from the first attributed passages for the source
+            chunk_resource_name = (
+                self.aqa_response_buffer.answer.grounding_attributions[
+                    0
+                ].source_id.semantic_retriever_chunk.chunk
+            )
+            get_chunk_response = self.retriever_service_client.get_chunk(
+                name=chunk_resource_name
+            )
+            metadata = get_chunk_response.custom_metadata
+            for m in metadata:
+                if m.key == "url":
+                    url = m.string_value
+        except:
+            url = "URL unknown"
+        return url
+
     # Print the prompt on the terminal for debugging
     def print_the_prompt(self, prompt):
         print("#########################################")