Update greg_funcs.py

greg_funcs.py

@@ -1,15 +1,21 @@
-
-
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
 from torch import tensor as torch_tensor
 from datasets import load_dataset
 
 from langchain.llms import OpenAI
 from langchain.docstore.document import Document
+from langchain.prompts import PromptTemplate
 from langchain.chains.question_answering import load_qa_chain
 from langchain.chains.qa_with_sources import load_qa_with_sources_chain
-from langchain.prompts import PromptTemplate
+from langchain import LLMMathChain, SQLDatabase, SQLDatabaseChain, LLMChain
+from langchain.agents import initialize_agent, Tool
+
+import sqlite3
+#import pandas as pd
+import json
 
+# database
+cxn = sqlite3.connect('./data/mbr.db')
 
 """# import models"""
 
@@ -62,9 +68,11 @@ def get_text_fmt(qry, passages = mypassages, doc_embedding=mycorpus_embeddings):
         prediction_text.append(result)
     return prediction_text
 
+"""# LLM based qa functions"""
+
 template = """You are a friendly AI assistant for the insurance company Humana. Given the following extracted parts of a long document and a question, create a succinct final answer. 
 If you don't know the answer, just say that you don't know. Don't try to make up an answer.
-If the question is not about Humana, politely inform them that you are tuned to only answer questions about Humana.
+If the question is not about Humana, politely inform the user that you are tuned to only answer questions about Humana.
 QUESTION: {question}
 =========
 {context}
@@ -74,18 +82,104 @@ PROMPT = PromptTemplate(template=template, input_variables=["context", "question
 
 chain_qa = load_qa_chain(OpenAI(temperature=0), chain_type="stuff", prompt=PROMPT)
 
+def get_text_fmt(qry, passages = mypassages, doc_embedding=mycorpus_embeddings):
+  predictions = search(qry, passages = passages, doc_embedding = doc_embedding, top_n=5, )
+  prediction_text = []
+  for hit in predictions:
+    page_content = passages[hit['corpus_id']]
+    metadata = {"source": hit['corpus_id']}
+    result = Document(page_content=page_content, metadata=metadata)
+    prediction_text.append(result)
+  return prediction_text
 
 def get_llm_response(message):
-    mydocs = get_text_fmt(message)
-    response = chain_qa.run(input_documents=mydocs, question=message)
-    return response
+  mydocs = get_text_fmt(message)
+  responses = chain_qa.run(input_documents=mydocs, question=message)
+  return responses
+
+# for x in xmpl_list:
+#   print(32*'=')
+#   print(x)
+#   print(32*'=')
+#   r = get_llm_response(x)
+#   print(r)
+
+"""# Database query"""
+
+db = SQLDatabase.from_uri("sqlite:///./data/mbr.db")
+
+llm = OpenAI(temperature=0)
+# default model
+# model_name: str = "text-davinci-003"
+# instruction fine-tuned, sometimes referred to as GPT-3.5
+
+db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)
+
+def db_qry(qry):
+  responses = db_chain.run(query='my mbr_id is 456 ;'+str(qry) ) ############### hardcode mbr id 456 for demo
+  return responses
+
+#db_qry('how many footcare visits have I had?')
+
+"""## Math
+- default version
+"""
+
+llm_math_chain = LLMMathChain(llm=llm, verbose=True)
+
+#llm_math_chain.run('what is the square root of 49?')
+
+"""# Greeting"""
+
+template = """You are a friendly AI assistant for the insurance company Humana. 
+Your name is Bruce and you were created on February 13, 20203.
+Offer polite greetings and brief small talk.
+Respond to thanks with, 'Glad to help.'
+If the question is not about Humana, politely guide the user to ask questions about Humana insurance benefits.
+QUESTION: {question}
+=========
+FINAL ANSWER:"""
+greet_prompt = PromptTemplate(template=template, input_variables=["question"])
+
+greet_llm = LLMChain(prompt=greet_prompt, llm=llm, verbose=True)
+
+"""# MRKL Chain"""
+
+tools = [
+    Tool(
+        name = "Benefit",
+        func=get_llm_response,
+        description="useful for when you need to answer questions about plan benefits, premiums and payments. You should ask targeted questions"
+    ),
+    Tool(
+        name="Calculator",
+        func=llm_math_chain.run,
+        description="useful for when you need to answer questions about math"
+    ),
+    Tool(
+        name="Member DB",
+        func=db_qry,
+        description="useful for when you need to answer questions about member details such their accumulated use of services. Input should be in the form of a question containing full context"
+    ),
+    Tool(
+        name="Greeting",
+        func=greet_llm.run,
+        description="useful for when you need to respond to greetings, thanks and make small talk"
+    ),
+]
+
+mrkl = initialize_agent(tools, llm, agent="zero-shot-react-description", verbose=True, return_intermediate_steps=True, max_iterations=5, early_stopping_method="generate")
+
+def mrkl_rspnd(qry):
+  response = mrkl({"input":str(qry) })
+  return response
 
 def chat(message, history):
     history = history or []
     message = message.lower()
     
-    response = get_llm_response(message)
-    history.append((message, response))
+    response = mrkl_rspnd(message)
+    history.append((message, response['output']))
     return history, history