diff --git a/lars/nepho/models/ask_sage_model.py b/lars/nepho/models/ask_sage_model.py index 58d9929..fbbdf11 100644 --- a/lars/nepho/models/ask_sage_model.py +++ b/lars/nepho/models/ask_sage_model.py @@ -1,5 +1,6 @@ import asyncio import json +import os from typing import List, Optional from .base_model import BaseModel from ..config import config @@ -10,8 +11,8 @@ class AskSageModel(BaseModel): """Ask Sage model implementation using the asksageclient API.""" - def __init__(self, model_name: str, credentials_json: str): - super().__init__(model_name) + def __init__(self, model_name: str, credentials_json: str, downscale_factor: Optional[int] = None): + super().__init__(model_name, downscale_factor=downscale_factor) self.credentials = _load_credentials(credentials_json) self.api_key = self.credentials['credentials']['api_key'] self.email = self.credentials['credentials']['Ask_sage_user_info']['username'] @@ -28,20 +29,33 @@ async def chat(self, prompt: str, images: Optional[List[str]] = None) -> str: loop = asyncio.get_event_loop() if images and self.supports_vision(): - for image_path in images: - if not self.validate_image(image_path): - raise ValueError(f"Invalid image: {image_path}") + temp_paths = [] + try: + prepared_images = [] + for image_path in images: + if not self.validate_image(image_path): + raise ValueError(f"Invalid image: {image_path}") - # query_with_file accepts a single path or a list - file_arg = images[0] if len(images) == 1 else images - response = await loop.run_in_executor( - None, - lambda: self.client.query_with_file( - message=prompt, - file=file_arg, - model=self.model_name + if self.downscale_factor and self.downscale_factor > 1: + prepared_path = self._downscale_image(image_path) + temp_paths.append(prepared_path) + else: + prepared_path = image_path + prepared_images.append(prepared_path) + + # query_with_file accepts a single path or a list + file_arg = prepared_images[0] if len(prepared_images) == 1 else prepared_images + response = await loop.run_in_executor( + None, + lambda: self.client.query_with_file( + message=prompt, + file=file_arg, + model=self.model_name + ) ) - ) + finally: + for temp_path in temp_paths: + os.remove(temp_path) else: response = await loop.run_in_executor( None, diff --git a/lars/nepho/models/base_model.py b/lars/nepho/models/base_model.py index a2c857d..25f8a62 100644 --- a/lars/nepho/models/base_model.py +++ b/lars/nepho/models/base_model.py @@ -1,6 +1,7 @@ import base64 import io import os +import tempfile from ..config import config from abc import ABC, abstractmethod @@ -10,20 +11,47 @@ class BaseModel(ABC): """Abstract base class for all chatbot models.""" - - def __init__(self, model_name: str): + + def __init__(self, model_name: str, downscale_factor: Optional[int] = None): + if downscale_factor is not None and (not isinstance(downscale_factor, int) or downscale_factor < 1): + raise ValueError("downscale_factor must be a positive integer") self.model_name = model_name - + self.downscale_factor = downscale_factor + @abstractmethod async def chat(self, prompt: str, images: Optional[List[str]] = None) -> str: """Generate a response based on the prompt and optional images.""" pass - + + def _downscale_image(self, image_path: str) -> str: + """Write a downscaled copy of the image to a temp file and return its path.""" + with Image.open(image_path) as img: + img_format = img.format or "PNG" + new_size = ( + max(1, img.width // self.downscale_factor), + max(1, img.height // self.downscale_factor), + ) + resized = img.resize(new_size, Image.LANCZOS) + suffix = os.path.splitext(image_path)[1] or ".png" + tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False) + try: + resized.save(tmp.name, format=img_format) + finally: + tmp.close() + return tmp.name + def encode_image(self, image_path: str) -> str: - """Encode image to base64 string for API calls.""" + """Encode image to base64 string for API calls, downscaling first if configured.""" try: - with open(image_path, "rb") as image_file: - return base64.b64encode(image_file.read()).decode('utf-8') + temp_path = None + if self.downscale_factor and self.downscale_factor > 1: + temp_path = self._downscale_image(image_path) + try: + with open(temp_path or image_path, "rb") as image_file: + return base64.b64encode(image_file.read()).decode('utf-8') + finally: + if temp_path: + os.remove(temp_path) except Exception as e: raise ValueError(f"Error encoding image {image_path}: {e}") diff --git a/lars/nepho/models/gpt_model.py b/lars/nepho/models/gpt_model.py index d8949d8..376ff33 100644 --- a/lars/nepho/models/gpt_model.py +++ b/lars/nepho/models/gpt_model.py @@ -7,9 +7,10 @@ class GPTModel(BaseModel): """GPT model implementation using OpenAI API.""" - def __init__(self, model_name: str = None, api_key: str = None, base_url: str = None, temperature: float = 0.7): + def __init__(self, model_name: str = None, api_key: str = None, base_url: str = None, temperature: float = 0.7, + downscale_factor: Optional[int] = None): model_name = model_name or config.DEFAULT_GPT_MODEL - super().__init__(model_name) + super().__init__(model_name, downscale_factor=downscale_factor) self.temperature = temperature self.api_key = api_key or config.OPENAI_API_KEY diff --git a/lars/nepho/models/ollama_model.py b/lars/nepho/models/ollama_model.py index 14e2396..d436b8c 100644 --- a/lars/nepho/models/ollama_model.py +++ b/lars/nepho/models/ollama_model.py @@ -8,9 +8,10 @@ class OllamaModel(BaseModel): """Ollama model implementation for local models.""" - def __init__(self, model_name: str = None, base_url: str = None, num_ctx: int = None): + def __init__(self, model_name: str = None, base_url: str = None, num_ctx: int = None, + downscale_factor: Optional[int] = None): model_name = model_name or config.DEFAULT_OLLAMA_MODEL - super().__init__(model_name) + super().__init__(model_name, downscale_factor=downscale_factor) self.base_url = base_url or config.OLLAMA_BASE_URL self.num_ctx = num_ctx or config.OLLAMA_NUM_CTX diff --git a/tests/test_base_model.py b/tests/test_base_model.py new file mode 100644 index 0000000..b4017b2 --- /dev/null +++ b/tests/test_base_model.py @@ -0,0 +1,45 @@ +import base64 +import io + +import pytest +from PIL import Image + +from lars.nepho.models.gpt_model import GPTModel + + +def _make_test_image(path, size=(40, 20)): + Image.new("RGB", size, color="red").save(path, format="PNG") + + +@pytest.fixture +def model(): + return GPTModel(model_name="gpt-4", api_key="test-key") + + +def test_downscale_factor_defaults_to_none(model): + assert model.downscale_factor is None + + +def test_invalid_downscale_factor_raises(): + with pytest.raises(ValueError, match="downscale_factor must be a positive integer"): + GPTModel(model_name="gpt-4", api_key="test-key", downscale_factor=0) + + +def test_encode_image_without_downscale_returns_original_bytes(model, tmp_path): + image_path = tmp_path / "test.png" + _make_test_image(image_path) + + encoded = model.encode_image(str(image_path)) + + assert base64.b64decode(encoded) == image_path.read_bytes() + + +def test_encode_image_downscales_by_integer_factor(tmp_path): + image_path = tmp_path / "test.png" + _make_test_image(image_path, size=(40, 20)) + + model = GPTModel(model_name="gpt-4", api_key="test-key", downscale_factor=4) + encoded = model.encode_image(str(image_path)) + + decoded_image = Image.open(io.BytesIO(base64.b64decode(encoded))) + assert decoded_image.size == (10, 5)