llm_service.py

from contextlib import contextmanager
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
import pandas as pd
from openai import OpenAI
from data_service import DataService

class LLMService(object):
  def __init__(self):
    self._openAIKey = None
    self._data_service = None
  
  @contextmanager
  def build(self):
    llm_service = None
    if self._openAIKey is None:
      raise ValueError("OPEN AI key was not provided and there is no default value.")
    if self._data_service is None:
      raise ValueError("To get the patient documents, a data service must be provided before building the LLMService.")
    try:
      llm_service = DefaultLLMService(self._openAIKey, self._data_service)
      yield llm_service
    finally:
      if llm_service is not None:
        llm_service.close()
  
  def close(self):
    raise Exception("Should not use the base class")

  def get_summary(self, patient: str) -> str:
    raise Exception("Should not use the base class")
    
  def answer_query(self, patient: str, query: str) -> str:
    raise Exception("Should not use the base class")

  def with_key(self, api_key: str):
    self._openAIKey = api_key
    return self
  
  def with_data_service(self, data_service: DataService):
    self._data_service = data_service
    return self

class DefaultLLMService(LLMService):
  def __init__(self, api_key: str, data_service: DataService):
    self._api_key = api_key
    self._data_service = data_service
    self._chatclient=OpenAI(api_key=self._api_key)
    #TODO decide on embedding model, using one provided in notebook
    self._embed_model = SentenceTransformer("all-MiniLM-L6-v2")
    self.build_chromadb()

  def close(self):
    #raise Exception("Not implemented")
    pass

  def get_chromadb(self,clear=0):
    client=chromadb.Client(Settings(
      persist_directory="./chroma_db"
    ))
    collection_name = "patient_data"
    # TODO: If clear and collection exists
    #if clear:
    #  client.delete_collection(collection_name)    
    return client.get_or_create_collection(name=collection_name)
    
  # TODO: It probably makes no difference, but the reason I chose to use a decorator @initialize in data_service is that I learned somewhere that it was better to put as little logic in a constructor as possible because (at least in other languages), errors in constructors made everything complicated, and potentially slow initialization logic made things difficult to ... I don't remember. Debug or parallelize or something. Maybe all of them. The trade-off is, if you forgot to decorate the method with @initialize, bad things.
  def build_chromadb(self):
    #TODO replace with cleaned data url or move to service inputs
    collection = self.get_chromadb(clear=1)
    texts = self._data_service.get_documents()
    metadatas = self._data_service.get_document_metadatas()
    # TODO: I'm looking at the docs and I'm wondering 1. if this is the default embedding function anyway (I think it's good to bring it out and see we can adjust it; I'm just pointing this out), and 2. whether it would be equivalent to provide self._embed_model.encode as the embedding function of the collection at configuration time.
    embeddings = self._embed_model.encode(texts).tolist()

    collection.add(
      documents=texts,
      embeddings=embeddings,
      metadatas=metadatas, #store everything except clinical notes and ids as metadata
      ids=[str(i) for i in range(len(texts))]
    )
  
  def query_chromadb(self, patient: str, query: str, result_template:str = "", top_n=3) -> str:
    if (patient=="") or (patient is None):
      return ""
    if (query=="") or (query is None):
      return ""

    collection = self.get_chromadb()
    query_embedding = self._embed_model.encode([query])[0].tolist()

    results = collection.query(
      query_embeddings=[query_embedding],
      n_results=top_n,
      # include=["documents","metadatas","distances"], #these are default query outputs so no need to specifiy
      where={"PATIENT_ID":patient} # specify patient
    )

    #TODO refine template for what info to include
    if result_template=="":
      result_template="""
      #{rank}: {desc} {st_date}\n
      {note}\n
      """

    context=""
    for i in range(len(results["ids"][0])):
      result_txt = result_template.format(
          rank=(i+1), #range is 0 indexed, increment for rank
          desc=results["metadatas"][0][i]["DESCRIPTION"],
          st_date=results["metadatas"][0][i]["START"],
          note=results["documents"][0][i])
      context = context+result_txt

    return context

  def get_summary(self, patient: str) -> str:
    #TODO get all visit notes or querying the vector database with specific prompt for a patient
    # all_visits = self._df.loc[self._df["PATIENT_ID"] == patient]
    vector_query=""
    vector_result_template="" #format for each result from vector search
    summary_context=self.query_chromadb(patient,vector_query,vector_result_template)
    summary_query=""
    # TODO: Add the found data to the context and asking OpenAI to summarize the docs provided?
    raise Exception("Not implemented")

  def answer_query(self, patient: str, query: str) -> str:
    # TODO: Example queries: tell me about incident ....
    #     Has this patient ...?
    #     Does this patient have a history of ...?
    # TODO: Find in vector database the most related docs to both 1. patient & 2. query
    rag=self.query_chromadb(patient,query)
    # TODO: Figure out how to utilize other columns.
    prompt_template="""
    You are an AI Assistant answering questions about a patient based on the relevant patient information provided.\n
    
    Patient Information:\n
    {RAG}

    Questions:\n
    {Query}

    Answer:"""
    filled_prompt=prompt_template.format(RAG=rag,Query=query)
    # Get model output
    response = self._chatclient.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "user", "content": filled_prompt}],
        temperature=0
    )
    # TODO: Error handling for 0 choices
    print(response)
    return response.choices[0].message.content