Merge pull request #185 from enoch3712/183-local-llm-llama32-model-not-working

test ollama fix

Author: enoch3712

tests/test_ollama.py

@@ -1,8 +1,6 @@
 import os
 from dotenv import load_dotenv
-from extract_thinker import LLM
-from extract_thinker import DocumentLoaderTesseract
-
+from extract_thinker.document_loader.document_loader_pypdf import DocumentLoaderPyPdf
 from extract_thinker.extractor import Extractor
 from tests.models.invoice import InvoiceContract
 
@@ -11,23 +9,20 @@
 
 
 def test_extract_with_ollama():
-
-    # Arrange
-    tesseract_path = os.getenv("TESSERACT_PATH")
-    test_file_path = os.path.join(cwd, "test_images", "invoice.png")
+    test_file_path = os.path.join(cwd, "tests", "files", "invoice.pdf")
 
     extractor = Extractor()
     extractor.load_document_loader(
-        DocumentLoaderTesseract(tesseract_path)
+        DocumentLoaderPyPdf()
     )
 
-    llm = LLM("ollama/phi3", "http://localhost:11434")
-    extractor.load_llm(llm)
+    os.environ["API_BASE"] = "http://localhost:11434"
+    extractor.load_llm("ollama/phi3.5")
 
     # Act
     result = extractor.extract(test_file_path, InvoiceContract)
 
     # Assert
     assert result is not None
-    assert result.invoice_number == "0000001"
-    assert result.invoice_date == "2014-05-07"
+    assert result.invoice_number == "00012"
+    assert result.invoice_date == "1/30/23"