test/test_llm_service.py

import unittest
from llm_service import LLMService
import os

class TestLlmService(unittest.TestCase):
  def test_initialize(self):
    """
    Tests that the service throws a ValueError when it is initialized without
    providing an API key, and then correctly initializes with one.
    """
    try:
      with LLMService().build():
        pass
      self.fail("This should have raised an exception.")
    except ValueError as e:
      pass

    # Written in a slightly strange way to test that the service initialized when given a key, but still allowed the close function to raise an exception if that hasn't been implemented yet.
    generator = LLMService().with_key("a key").build()
    generator.__enter__()
    try:
      temp3.__exit__(None, None, None)
    except Exception as e:
      pass

  def test_simple_query(self):
    """
    Tests querying the chroma database.
    """
    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)