Working answer query

llm_service.py

@@ -3,6 +3,7 @@ import chromadb
 from chromadb.config import Settings
 from sentence_transformers import SentenceTransformer
 import pandas as pd
+from openai import OpenAI
 
 class LLMService(object):
   def __init__(self):
@@ -17,7 +18,8 @@ class LLMService(object):
       llm_service = DefaultLLMService(self._openAIKey)
       yield llm_service
     finally:
-      llm_service.close()
+      if llm_service is not None:
+        llm_service.close()
   
   def close(self):
     raise Exception("Should not use the base class")
@@ -41,14 +43,16 @@ class DefaultLLMService(LLMService):
     self.build_chromadb()
 
   def close(self):
-    raise Exception("Not implemented")
+    #raise Exception("Not implemented")
+    pass
 
   def get_chromadb(self,clear=0):
     client=chromadb.Client(Settings(
       persist_directory="./chroma_db"
     ))
-    if clear:
-      client.delete_collection(collection_name)    
+    collection_name = "patient_data"
+    #if clear:
+    #  client.delete_collection(collection_name)    
     return client.get_or_create_collection(name=collection_name)
     
   def build_chromadb(self):
@@ -62,7 +66,7 @@ class DefaultLLMService(LLMService):
     collection.add(
       documents=texts,
       embeddings=embeddings,
-      metadatas=[df.iloc[i,3:11].to_dict() for i in range(len(texts))], #store everything except clinical notes and ids as metadata
+      metadatas=[self._df.iloc[i,3:11].to_dict() for i in range(len(texts))], #store everything except clinical notes and ids as metadata
       ids=[str(i) for i in range(len(texts))]
     )
   
@@ -115,7 +119,7 @@ class DefaultLLMService(LLMService):
     #     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)
+    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
@@ -134,4 +138,6 @@ class DefaultLLMService(LLMService):
         messages=[{"role": "user", "content": filled_prompt}],
         temperature=0
     )
-    return response
+    # TODO: Error handling for 0 choices
+    print(response)
+    return response.choices[0].message.content

requirements.txt

@@ -4,4 +4,5 @@ openai
 python-dotenv
 langchain_openai
 chromadb
-sentence-transformers
+sentence-transformers
+numpy<2

test/test_llm_service.py

@@ -1,5 +1,6 @@
 import unittest
 from llm_service import LLMService
+import os
 
 class TestLlmService(unittest.TestCase):
   def test_initialize(self):
@@ -22,8 +23,26 @@ class TestLlmService(unittest.TestCase):
     except Exception as e:
       pass
 
-  def test_get_summary(self):
+  def test_simple_query(self):
     """
-    Tests the generation of the patient summary.
+    Tests querying the chroma database.
     """
-    self.fail("Test not implemented yet")
+    with LLMService().with_key(os.getenv("OPENAI_API_KEY")).build() as llm_service:
+      result = llm_service.query_chromadb(0, "tell me about dental visits", "")
+      print(result)
+
+  def test_simple_summary(self):
+    """
+    Tests the patient summary
+    """
+    with LLMService().with_key(os.getenv("OPENAI_API_KEY")).build() as llm_service:
+      result = llm_service.get_summary(0)
+      print(result)
+
+  def test_simple_answer_query(self):
+    """
+    Tests a simple query
+    """
+    with LLMService().with_key(os.getenv("OPENAI_API_KEY")).build() as llm_service:
+      result = llm_service.answer_query(0, "tell me about dental visits")
+      print(result)