Upload herbal_expert.py

herbal_expert.py

@@ -16,11 +16,12 @@ import nltk
 from langchain.callbacks.manager import CallbackManagerForLLMRun
 from typing import List, Union, Optional, Any
 
-ngrok_url = 'https://2590-2605-7b80-3d-320-a515-4f0d-f60e-71e5.ngrok-free.app/'
-ngrok_url = 'http://localhost:1234/'
+ngrok_url = 'https://9c1a-2605-7b80-3d-320-fc74-5877-9733-e99b.ngrok-free.app/'
+#ngrok_url = 'http://localhost:1234/'
 Entrez.email = "[email protected]"
 nltk.download('wordnet')
 
+
 class CustomLLM(LLM):
     n: int
 
@@ -35,6 +36,18 @@ class CustomLLM(LLM):
             run_manager: Optional[CallbackManagerForLLMRun] = None,
             **kwargs: Any,
     ) -> str:
+        """
+        The _call function is the function that will be called by the user.
+        It should take in a prompt and return a response.
+
+        :param self: Represent the instance of the class
+        :param prompt: str: Pass the prompt to the model
+        :param stop: Optional[List[str]]: Define the stop words
+        :param run_manager: Optional[CallbackManagerForLLMRun]: Pass the run manager to the call function
+        :param **kwargs: Any: Pass in any additional parameters that may be needed for the function
+        :param : Pass the prompt to the model
+        :return: A string that is the response of gpt-3 to the prompt
+        """
         data = {
             "messages": [
                 {
@@ -50,7 +63,6 @@ class CustomLLM(LLM):
         return json.loads(response.text)['choices'][0]['message']['content']
 
 
-
 class CustomPromptTemplate(StringPromptTemplate):
     template: str
 
@@ -82,6 +94,14 @@ prompt_with_history = CustomPromptTemplate(
 
 
 def get_agent_executor():
+    """
+    The get_agent_executor function is a factory function that returns an AgentExecutor object.
+    The AgentExecutor object is the main interface for interacting with the agent.
+    agent.run('query') is the main method for interacting with the agent.
+
+    :return: An AgentExecutor object
+    """
+
     llm = CustomLLM(n=10)
     question_decompose_chain = LLMChain(llm=llm, prompt=prompt_with_history)
     question_decompose_agent = LLMSingleActionAgent(
@@ -97,15 +117,14 @@ def get_agent_executor():
         tools=[],
         verbose=True,
         memory=memory
-        )
+    )
 
 
 def get_num_citations(pmid: str):
     """
     The get_num_citations function takes a PubMed ID (pmid) as input and returns the number of citations for that
-    pmid. The function uses the Entrez module to query PubMed Central's API, which is part of NCBI's Entrez
-    Programming Utilities. The function first queries PMC using elink to get all articles citing the given pmid,
-    then it counts how many articles are in that list.
+    pmid. The function uses the Entrez module to query PubMed Central's API. The function first queries PMC using
+    elink to get all articles citing the given pmid, then it counts how many articles are in that list.
 
     :param pmid: str: Specify the pmid of the article you want to get citations for
     :return: The number of citations for a given pmid
@@ -120,16 +139,18 @@ def get_num_citations(pmid: str):
         else:
             return 0
 
+
 def fetch_pubmed_articles(keywords, max_search=10, max_context=3):
     """
-    The fetch_pubmed_articles function takes in a list of keywords and returns a list of articles.
-        The function uses the Entrez API to search for articles with the given keywords, then fetches
-        those articles from PubMed. The function returns a list of strings, where each string is an article.
-
-    :param keywords: Search for articles in the pubmed database
-    :param max_results: Specify the number of articles to be returned default is 1
-    :param email: Identify the user to ncbi
-    :return: A list of strings
+    The fetch_pubmed_articles function takes in a list of keywords and returns the top 3 articles from PubMed that
+    are most relevant to those keywords. First the search is done on max_search articles, the list is then sorted by
+    number of citations, then the top max_content articles are chosen from that list. If no articles are found with
+    the initial list of keywords, the search is rerun with the top 4 keywords of the list
+
+    :param keywords: Search for articles in pubmed
+    :param max_search: Limit the number of initial search results
+    :param max_context: Specify the number of articles to return
+    :return: A list of articles
     """
 
     try:
@@ -178,11 +199,10 @@ def call_model_with_history(messages: list):
         "stop": ["### Instruction:"], "temperature": 0, "max_tokens": 512, "stream": False
     }
 
-    response = requests.post(ngrok_url+ "v1/chat/completions", headers={"Content-Type": "application/json"}, json=data)
+    response = requests.post(ngrok_url + "v1/chat/completions", headers={"Content-Type": "application/json"}, json=data)
     return json.loads(response.text)['choices'][0]['message']['content']
 
 
-
 # TODO: add ability to pass message history to model
 def format_prompt_and_query(prompt, **kwargs):
     """
@@ -190,7 +210,7 @@ def format_prompt_and_query(prompt, **kwargs):
     arguments, and then calls call_model_with_history with a list of messages containing the formatted prompt.
 
     :param prompt: Format the prompt with the values in kwargs
-    :param **kwargs: Pass a dictionary of key-value pairs to the formatting function
+    :param **kwargs: Pass a dictionary of key-value pairs to the prompt formatting function
     :return: A list of dictionaries
     """
 
@@ -221,37 +241,39 @@ class HerbalExpert:
         ]
         # og = Original, qa = Question Asking, ri = Response Improvement
         self.prompts = {
-            "og_answer_prompt": """### Instruction: Answer the following question to the best of your ability. Question: {question} 
-                    ### Response: Answer: """,
+            "og_answer_prompt": """### Instruction: Answer the following question to the best of your ability. 
+                Question: {question} 
+                ### Response: Answer: """,
 
             "ans_decompose_prompt": """### Instruction: Given the following text, identify the 2 most important 
-                    keywords that capture the essence of the text. If there's a list of products, choose the top 2 products. 
-                    Your response should be a list of only 2 keywords separated by commas. 
-                    Text: {original_answer} 
-                    ### Response: Keywords: """,
+                keywords that capture the essence of the text. If there's a list of products, choose the top 2 products. 
+                Your response should be a list of only 2 keywords separated by commas. 
+                Text: {original_answer} 
+                ### Response: Keywords: """,
 
             "qa_prompt": """### Instruction: Answer the following question using the given context ONLY if the 
-                    context is relevant to the question. If the context doesn't help answer the question respond with "I don't know". 
-                    Question: {question} 
-                    Context: {context} 
-                    ### Response: Answer: """,
+                context is relevant to the question. If the context doesn't help answer the question respond with "I don't know". 
+                Question: {question} 
+                Context: {context} 
+                ### Response: Answer: """,
 
             "ri_prompt": """### Instruction: You are an caring, intelligent question answering agent. Craft a 
-                    response that is more informative and intelligent than the original answer and imparts knowledge from 
-                    both the old answer and from the context ONLY if it helps answer the question. 
-                    Question: {question} 
-                    Old Answer: {answer} 
-                    Context: {answer2} 
-                    ### Response: Improved Answer: """
+                response that is more informative and intelligent than the original answer and imparts knowledge from 
+                both the old answer and from the context ONLY if it helps answer the question. 
+                Question: {question} 
+                Old Answer: {answer} 
+                Context: {answer2} 
+                ### Response: Improved Answer: """
         }
 
     def process_query_words(self, question_words: str, answer_words: str):
         # don't need to be searching for these in pubmed. Should we include: 'supplements', 'supplement'
         """
         The process_query_words function takes in a string of words and returns a list of filtered lemmatized words.
-        The function first splits the input strings into lists, then removes any duplicate entries from the list.
-        It then iterates through each word in the list and strips it of whitespace before passing it to WordNetLemmatizer().lemmatize()
-        to return its lemma (base form). The function also removes any vague words that are not useful for searching PubMed.
+        The function first splits the input strings into lists, then removes any duplicate entries from the list. It
+        then iterates through each word in the list and strips it of whitespace before passing it to
+        WordNetLemmatizer().lemmatize() to return its lemma (base form). The function also removes any vague words
+        that are not useful for searching PubMed.
 
         :param self: Represent the instance of the class
         :param question_words: str: Get the question words from the user
@@ -297,6 +319,19 @@ class HerbalExpert:
         return words, original_answer
 
     def query_expert(self, question: str = None):
+        """
+        The query_expert function takes a question as input and returns the expert's response to that question. The
+        function first converts the question into keywords, then uses those keywords to search PubMed for relevant
+        articles. If no articles are found, it returns the original response from the expert (i.e., without context).
+        If at least one article is found, it asks two follow-up questions: 1) "What do you think of this answer
+        in light of these new findings?" and 2) "How would you improve your answer based on these new
+        findings?". It then returns both responses.
+
+        :param self: Represent the instance of the class
+        :param question: str: Pass in the question that is to be asked
+        :return: A dictionary with the question, response and run info
+        """
+
         question = random.choice(self.default_questions) if question is None else question
         print("Question: ", question)
 
@@ -326,9 +361,9 @@ class HerbalExpert:
 
 herbal_expert = HerbalExpert(get_agent_executor())
 
-
 if __name__ == '__main__':
     herbal_expert = HerbalExpert(get_agent_executor())
-    answer = herbal_expert.query_expert("I'm experiencing consistent stress and anxiety. What herbs or supplements could help alleviate these symptoms?")
+    answer = herbal_expert.query_expert(
+        "I'm experiencing consistent stress and anxiety. What herbs or supplements could help alleviate these symptoms?")
     print(answer['response'])
     # return to api? who knows