Merge pull request #18 from enoch3712/17-ollama-support-for-local-llms

Ollama support added. With a test connecting to a phi-3 mini model

Author: enoch3712

extract_thinker/llm.py

@@ -6,10 +6,13 @@
 
 
 class LLM:
-    def __init__(self, model: str):
+    def __init__(self, model: str, api_base: str = None, api_key: str = None, api_version: str = None):
         self.client = instructor.from_litellm(litellm.completion, mode=instructor.Mode.MD_JSON)
         self.model = model
         self.router = None
+        self.api_base = api_base
+        self.api_key = api_key
+        self.api_version = api_version
 
     def load_router(self, router: Router) -> None:
         self.router = router
@@ -31,7 +34,10 @@ def request(self, messages: List[Dict[str, str]], response_model: str) -> Any:
                 model=self.model,
                 max_tokens=max_tokens,
                 messages=messages,
-                response_model=response_model
+                response_model=response_model,
+                api_base=self.api_base,
+                api_key=self.api_key,
+                api_version=self.api_version
             )
 
         return response

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "extract_thinker"
-version = "0.0.5"
+version = "0.0.6"
 description = "Library to extract data from files and documents agnositicaly using LLMs"
 authors = ["Júlio Almeida <[email protected]>"]
 readme = "README.md"