diff --git a/.env.template b/.env.template index 47b7b6f..2e798ac 100644 --- a/.env.template +++ b/.env.template @@ -1,3 +1,4 @@ # LLM API Keys ANTHROPIC_API_KEY=sk-ant-your-key-here OPENAI_API_KEY=sk-your-key-here +OLLAMA_API_KEY=you-key-here diff --git a/example.py b/example.py index 069511d..a613eaf 100644 --- a/example.py +++ b/example.py @@ -5,6 +5,7 @@ from harvestor import Harvestor, list_models + load_dotenv() @@ -18,24 +19,33 @@ class SimpleInvoiceSchema(BaseModel): None, description="The customer firstname" ) customer_lastname: Optional[str] = Field(None, description="The customer lastname") + invoice_total_price_with_taxes: Optional[float] = Field( + None, description="The total price with taxes" + ) + invoice_total_price_without_taxes: Optional[float] = Field( + None, description="The total price without taxes" + ) # List available models print("Available models:", list(list_models().keys())) # Use default model (claude-haiku) -h = Harvestor(model="claude-haiku") +# h = Harvestor(model="claude-haiku") -output = h.harvest_file( - source="data/uploads/keep_for_test.jpg", schema=SimpleInvoiceSchema -) +# output = h.harvest_file( +# source="data/uploads/keep_for_test.jpg", schema=SimpleInvoiceSchema +# ) -print(output.to_summary()) +# print(output.to_summary()) # Alternative: use OpenAI # h_openai = Harvestor(model="gpt-4o-mini") -# output = h_openai.harvest_file("invoice.jpg", schema=SimpleInvoiceSchema) +# output = h_openai.harvest_file("data/uploads/keep_for_test.jpg", schema=SimpleInvoiceSchema) -# Alternative: use local Ollama (free) -# h_ollama = Harvestor(model="llava") -# output = h_ollama.harvest_file("invoice.jpg", schema=SimpleInvoiceSchema) +# Alternative: use local Ollama (free) or cloud Ollama +h_ollama = Harvestor(model="gemma3:4b-cloud") +output = h_ollama.harvest_file( + "data/uploads/keep_for_test.jpg", schema=SimpleInvoiceSchema +) +print(output.to_summary()) diff --git a/pyproject.toml b/pyproject.toml index d72ebe4..42adda9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "langchain-openai>=0.0.5", "anthropic>=0.18.0", "openai>=1.10.0", - "httpx>=0.27.0", # For Ollama provider + "ollama>=0.6.1", # Document Processing "PyMuPDF>=1.23.0", diff --git a/src/harvestor/providers/ollama.py b/src/harvestor/providers/ollama.py index 1bc03d5..fa866a5 100644 --- a/src/harvestor/providers/ollama.py +++ b/src/harvestor/providers/ollama.py @@ -6,7 +6,7 @@ import os from typing import Optional -import httpx +from ollama import generate, Client, list as list_models from .base import BaseLLMProvider, CompletionResult, ModelInfo @@ -66,17 +66,25 @@ def __init__( if model not in OLLAMA_MODELS: # Allow custom models not in the predefined list self.model_config = { - "id": f"{model}:latest" if ":" not in model else model, + "id": f"{model}:latest" + if ":" not in model + else model, # allow models like ministral-3:3b "input_cost": 0.0, "output_cost": 0.0, - "supports_vision": False, + "supports_vision": True, "context_window": 8192, } else: self.model_config = OLLAMA_MODELS[model] + self.client = None # if self.client -> using ollama cloud + if model.endswith("cloud"): + self.client = Client( + host="https://ollama.com", + headers={"Authorization": "Bearer " + os.environ.get("OLLAMA_API_KEY")}, + ) + self.model_id = self.model_config["id"] - self.client = httpx.Client(base_url=base_url, timeout=120.0) def complete( self, @@ -85,20 +93,12 @@ def complete( temperature: float = 0.0, ) -> CompletionResult: try: - response = self.client.post( - "/api/generate", - json={ - "model": self.model_id, - "prompt": prompt, - "stream": False, - "options": { - "temperature": temperature, - "num_predict": max_tokens, - }, - }, + data = generate( + model=self.model_id, + prompt=prompt, + stream=False, + options={"temperature": temperature, "num_predict": max_tokens}, ) - response.raise_for_status() - data = response.json() return CompletionResult( success=True, @@ -112,13 +112,6 @@ def complete( }, ) - except httpx.ConnectError: - return CompletionResult( - success=False, - content="", - model=self.model_id, - error=f"Cannot connect to Ollama at {self.base_url}. Is Ollama running?", - ) except Exception as e: return CompletionResult( success=False, @@ -146,27 +139,32 @@ def complete_vision( try: image_b64 = base64.standard_b64encode(image_data).decode("utf-8") - response = self.client.post( - "/api/generate", - json={ - "model": self.model_id, - "prompt": prompt, - "images": [image_b64], - "stream": False, - "options": { - "temperature": temperature, - "num_predict": max_tokens, - }, - }, - ) - response.raise_for_status() - data = response.json() + if self.client: + data = self.client.generate( + model=self.model, + prompt=prompt, + images=[image_b64], + stream=False, + options={"temperature": temperature, "num_predict": max_tokens}, + ) + else: + data = generate( + model=self.model, + prompt=prompt, + images=[image_b64], + stream=False, + options={"temperature": temperature, "num_predict": max_tokens}, + ) return CompletionResult( success=True, content=data.get("response", ""), - input_tokens=data.get("prompt_eval_count", 0), - output_tokens=data.get("eval_count", 0), + input_tokens=0 + if data.get("prompt_eval_count") is None + else data["prompt_eval_count"], + output_tokens=0 + if data.get("eval_count") is None + else data["eval_count"], model=self.model_id, metadata={ "total_duration": data.get("total_duration"), @@ -174,13 +172,6 @@ def complete_vision( }, ) - except httpx.ConnectError: - return CompletionResult( - success=False, - content="", - model=self.model_id, - error=f"Cannot connect to Ollama at {self.base_url}. Is Ollama running?", - ) except Exception as e: return CompletionResult( success=False, @@ -210,9 +201,7 @@ def get_provider_name(cls) -> str: def list_local_models(self) -> list[str]: """List models available in the local Ollama installation.""" try: - response = self.client.get("/api/tags") - response.raise_for_status() - data = response.json() + data = list_models() return [m["name"] for m in data.get("models", [])] except Exception: return [] diff --git a/src/harvestor/schemas/base.py b/src/harvestor/schemas/base.py index 3232344..61b0d6f 100644 --- a/src/harvestor/schemas/base.py +++ b/src/harvestor/schemas/base.py @@ -188,7 +188,7 @@ def to_summary(self) -> str: status = "SUCCESS" if self.success else "FAILED" strategy = self.final_strategy.value if self.final_strategy else "N/A" - return f""" + summary = f""" Harvest Result: {status} Document: {self.document_id} ({self.document_type}) Strategy: {strategy} @@ -199,6 +199,13 @@ def to_summary(self) -> str: Data: {self.data} """.strip() + if self.error: + summary += f""" + Error: {self.error} + """ + + return summary + @dataclass class CostReport: diff --git a/uv.lock b/uv.lock index 9f84b22..9e2e68f 100644 --- a/uv.lock +++ b/uv.lock @@ -271,10 +271,10 @@ source = { editable = "." } dependencies = [ { name = "anthropic" }, { name = "click" }, - { name = "httpx" }, { name = "langchain" }, { name = "langchain-anthropic" }, { name = "langchain-openai" }, + { name = "ollama" }, { name = "openai" }, { name = "opencv-python-headless" }, { name = "pdfplumber" }, @@ -306,10 +306,10 @@ dev = [ requires-dist = [ { name = "anthropic", specifier = ">=0.18.0" }, { name = "click", specifier = ">=8.1.0" }, - { name = "httpx", specifier = ">=0.27.0" }, { name = "langchain", specifier = ">=0.1.0" }, { name = "langchain-anthropic", specifier = ">=0.1.0" }, { name = "langchain-openai", specifier = ">=0.0.5" }, + { name = "ollama", specifier = ">=0.6.1" }, { name = "openai", specifier = ">=1.10.0" }, { name = "opencv-python-headless", specifier = ">=4.8.0" }, { name = "pdfplumber", specifier = ">=0.10.0" }, @@ -635,6 +635,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/67/0e/35082d13c09c02c011cf21570543d202ad929d961c02a147493cb0c2bdf5/numpy-2.2.6-cp313-cp313t-win_amd64.whl", hash = "sha256:6031dd6dfecc0cf9f668681a37648373bddd6421fff6c66ec1624eed0180ee06", size = 12771374 }, ] +[[package]] +name = "ollama" +version = "0.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/5a/652dac4b7affc2b37b95386f8ae78f22808af09d720689e3d7a86b6ed98e/ollama-0.6.1.tar.gz", hash = "sha256:478c67546836430034b415ed64fa890fd3d1ff91781a9d548b3325274e69d7c6", size = 51620 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/47/4f/4a617ee93d8208d2bcf26b2d8b9402ceaed03e3853c754940e2290fed063/ollama-0.6.1-py3-none-any.whl", hash = "sha256:fc4c984b345735c5486faeee67d8a265214a31cbb828167782dc642ce0a2bf8c", size = 14354 }, +] + [[package]] name = "openai" version = "2.16.0"