donkit-llm 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

donkit/llm/ollama_integration.py

@@ -0,0 +1,442 @@
+#!/usr/bin/env python3
+"""Test script for Ollama LLM provider integration."""
+
+import asyncio
+import json
+
+import pytest
+
+from llm_gate.models import ModelCapability
+from .factory import ModelFactory
+from .model_abstract import (
+    ContentPart,
+    ContentType,
+    EmbeddingRequest,
+    GenerateRequest,
+    Message,
+)
+from .openai_model import OpenAIModel
+
+
+def test_ollama_model_creation() -> None:
+    """Test that Ollama model initializes correctly."""
+    print("=" * 60)
+    print("TEST 1: Ollama Model Creation via Factory")
+    print("=" * 60)
+
+    credentials = {
+        "api_key": "ollama",
+        "base_url": "http://localhost:11434/v1",
+    }
+
+    try:
+        model = ModelFactory.create_model(
+            provider="ollama",
+            model_name="gpt-oss:20b-cloud",
+            credentials=credentials,
+        )
+        print(f"✅ Model created: {model.__class__.__name__}")
+        print(f"   Model name: {model.model_name}")
+        print(f"   Is OpenAI model: {isinstance(model, OpenAIModel)}")
+    except Exception as e:
+        print(f"❌ Failed to create model: {e}")
+        raise
+
+
+def test_openai_model_direct_creation() -> None:
+    """Test direct OpenAI model creation for Ollama."""
+    print("\n" + "=" * 60)
+    print("TEST 2: Direct OpenAI Model Creation for Ollama")
+    print("=" * 60)
+
+    try:
+        model = ModelFactory.create_openai_model(
+            model_name="gpt-oss:20b-cloud",
+            api_key="ollama",
+            base_url="http://localhost:11434/v1",
+        )
+        print(f"✅ Model created: {model.__class__.__name__}")
+        print(f"   Model name: {model.model_name}")
+        print(f"   Has client: {hasattr(model, 'client')}")
+        print(f"   Client type: {type(model.client).__name__}")
+    except Exception as e:
+        print(f"❌ Failed to create model: {e}")
+        raise
+
+
+def test_embedding_model_creation() -> None:
+    """Test embedding model creation for Ollama."""
+    print("\n" + "=" * 60)
+    print("TEST 3: Embedding Model Creation")
+    print("=" * 60)
+
+    try:
+        # OpenAI embedding model with Ollama base URL
+        model = ModelFactory.create_embedding_model(
+            provider="openai",
+            model_name="nomic-embed-text",
+            api_key="ollama",
+            base_url="http://localhost:11434/v1",
+        )
+        print(f"✅ Embedding model created: {model.__class__.__name__}")
+        print(f"   Model name: {model.model_name}")
+    except Exception as e:
+        print(f"❌ Failed to create embedding model: {e}")
+        raise
+
+
+def test_model_configuration() -> None:
+    """Test model configuration."""
+    print("\n" + "=" * 60)
+    print("TEST 4: Model Configuration")
+    print("=" * 60)
+
+    try:
+        model = ModelFactory.create_openai_model(
+            model_name="gpt-oss:20b-cloud",
+            api_key="ollama",
+            base_url="http://localhost:11434/v1",
+        )
+
+        config = {
+            "model_name": model.model_name,
+            "model_type": model.__class__.__name__,
+            "has_client": hasattr(model, "client"),
+            "capabilities": str(model.capabilities),
+        }
+
+        print("✅ Model configuration:")
+        print(json.dumps(config, indent=2, default=str))
+    except Exception as e:
+        print(f"❌ Failed to get model configuration: {e}")
+        raise
+
+
+def test_factory_provider_support() -> None:
+    """Test that factory supports ollama provider."""
+    print("\n" + "=" * 60)
+    print("TEST 5: Factory Provider Support")
+    print("=" * 60)
+
+    credentials = {
+        "api_key": "ollama",
+        "base_url": "http://localhost:11434/v1",
+    }
+
+    try:
+        # Test all supported providers
+        providers = ["openai", "ollama"]
+        for provider in providers:
+            model = ModelFactory.create_model(
+                provider=provider,
+                model_name="test-model",
+                credentials=credentials,
+            )
+            print(f"✅ Provider '{provider}' supported: {model.__class__.__name__}")
+    except Exception as e:
+        print(f"❌ Provider support test failed: {e}")
+        raise
+
+
+@pytest.mark.asyncio
+async def test_text_generation() -> None:
+    """Test text generation with Ollama."""
+    print("\n" + "=" * 60)
+    print("TEST 6: Text Generation Request")
+    print("=" * 60)
+
+    credentials = {
+        "api_key": "ollama",
+        "base_url": "http://localhost:11434/v1",
+    }
+
+    try:
+        model = ModelFactory.create_model(
+            provider="ollama",
+            model_name="gpt-oss:20b-cloud",
+            credentials=credentials,
+        )
+
+        request = GenerateRequest(
+            messages=[
+                Message(
+                    role="user", content="Say 'Hello from Ollama!' in one sentence."
+                )
+            ],
+            temperature=0.7,
+            max_tokens=100,
+        )
+
+        print("📤 Sending request to Ollama...")
+        print("   Model: gpt-oss:20b-cloud")
+        print(f"   Message: {request.messages[0].content}")
+
+        response = await model.generate(request)
+
+        print("✅ Generation successful")
+        print(f"   Response: {response.content}")
+
+    except Exception as e:
+        print(f"❌ Generation failed: {e}")
+        raise
+
+
+@pytest.mark.asyncio
+async def test_streaming_generation() -> None:
+    """Test streaming text generation with Ollama."""
+    print("\n" + "=" * 60)
+    print("TEST 7: Streaming Text Generation")
+    print("=" * 60)
+
+    credentials = {
+        "api_key": "ollama",
+        "base_url": "http://localhost:11434/v1",
+    }
+
+    try:
+        model = ModelFactory.create_model(
+            provider="ollama",
+            model_name="gpt-oss:20b-cloud",
+            credentials=credentials,
+        )
+
+        request = GenerateRequest(
+            messages=[
+                Message(role="user", content="Count from 1 to 5, one number per line.")
+            ],
+            temperature=0.5,
+            max_tokens=100,
+        )
+
+        print("📤 Sending streaming request to Ollama...")
+        print("   Model: gpt-oss:20b-cloud")
+        print(f"   Message: {request.messages[0].content}")
+        print("\n   Response stream:")
+
+        full_content = ""
+        async for chunk in model.generate_stream(request):
+            if chunk.content:
+                print(f"   {chunk.content}", end="", flush=True)
+                full_content += chunk.content
+
+        print("\n\n✅ Streaming complete")
+        print(f"   Total content length: {len(full_content)} characters")
+
+    except Exception as e:
+        print(f"❌ Streaming failed: {e}")
+        raise
+
+
+@pytest.mark.asyncio
+async def test_embedding_generation() -> None:
+    """Test embedding generation with Ollama."""
+    print("\n" + "=" * 60)
+    print("TEST 8: Embedding Generation")
+    print("=" * 60)
+    try:
+        model = ModelFactory.create_embedding_model(
+            provider="openai",
+            model_name="embeddinggemma",
+            api_key="ollama",
+            base_url="http://localhost:11434/v1",
+        )
+
+        texts = [
+            "Hello from Ollama!",
+            "This is a test embedding.",
+            "Embedding models are useful for semantic search.",
+        ]
+
+        request = EmbeddingRequest(input=texts)
+
+        print("📤 Sending embedding request to Ollama...")
+        print("   Model: nomic-embed-text")
+        print(f"   Texts to embed: {len(texts)}")
+        for i, text in enumerate(texts, 1):
+            print(f"     {i}. {text}")
+
+        response = await model.embed(request)
+
+        print("✅ Embedding successful")
+        print(f"   Number of embeddings: {len(response.embeddings)}")
+        print(f"   Embedding dimension: {len(response.embeddings[0])}")
+        print(f"   First embedding (first 5 values): {response.embeddings[0][:5]}")
+
+    except Exception as e:
+        print(f"❌ Embedding failed: {e}")
+        raise
+
+
+@pytest.mark.asyncio
+async def test_multimodal_vision() -> None:
+    """Test multimodal vision capabilities with Ollama."""
+    print("\n" + "=" * 60)
+    print("TEST 9: Multimodal Vision (Image Understanding)")
+    print("=" * 60)
+
+    credentials = {
+        "api_key": "ollama",
+        "base_url": "http://localhost:11434/v1",
+    }
+
+    try:
+        model = ModelFactory.create_model(
+            provider="ollama",
+            model_name="qwen3-vl:235b-cloud",
+            credentials=credentials,
+        )
+
+        # Check if model supports vision
+        print("📋 Model capabilities:")
+        print(f"   Vision support: {model.supports_capability(ModelCapability.VISION)}")
+        print(
+            f"   Multimodal input: {model.supports_capability(ModelCapability.MULTIMODAL_INPUT)}"
+        )
+
+        # Create a message with image URL (using a public test image)
+        image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg"
+
+        content_parts = [
+            ContentPart(
+                content_type=ContentType.TEXT,
+                content="What do you see in this image? Describe it briefly.",
+            ),
+            ContentPart(
+                content_type=ContentType.IMAGE_URL,
+                content=image_url,
+                mime_type="image/jpeg",
+            ),
+        ]
+
+        message = Message(role="user", content=content_parts)
+
+        request = GenerateRequest(
+            messages=[message],
+            temperature=0.7,
+            max_tokens=200,
+        )
+
+        print("\n📤 Sending multimodal request to Ollama...")
+        print("   Model: gpt-oss:20b-cloud")
+        print(f"   Message parts: {len(content_parts)}")
+        print("     - Text: What do you see in this image?")
+        print(f"     - Image: {image_url}")
+
+        response = await model.generate(request)
+
+        print("\n✅ Vision analysis successful")
+        print(f"   Response: {response.content}")
+
+    except Exception as e:
+        print(f"❌ Vision test failed: {e}")
+        raise
+
+
+@pytest.mark.asyncio
+async def test_multimodal_with_base64() -> None:
+    """Test multimodal with base64 encoded image."""
+    print("\n" + "=" * 60)
+    print("TEST 10: Multimodal with Base64 Image")
+    print("=" * 60)
+
+    credentials = {
+        "api_key": "ollama",
+        "base_url": "http://localhost:11434/v1",
+    }
+
+    try:
+        model = ModelFactory.create_model(
+            provider="ollama",
+            model_name="qwen3-vl:235b-cloud",
+            credentials=credentials,
+        )
+
+        # For this test, we'll use a simple base64 encoded 1x1 pixel image
+        # In real usage, you would encode an actual image
+        base64_image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+
+        content_parts = [
+            ContentPart(
+                content_type=ContentType.TEXT,
+                content="Analyze this image and tell me what you see.",
+            ),
+            ContentPart(
+                content_type=ContentType.IMAGE_BASE64,
+                content=base64_image,
+                mime_type="image/png",
+            ),
+        ]
+
+        message = Message(role="user", content=content_parts)
+
+        request = GenerateRequest(
+            messages=[message],
+            temperature=0.5,
+            max_tokens=100,
+        )
+
+        print("📤 Sending base64 image request to Ollama...")
+        print("   Model: gpt-oss:20b-cloud")
+        print(f"   Message parts: {len(content_parts)}")
+        print("     - Text: Analyze this image")
+        print(f"     - Image (base64): {len(base64_image)} characters")
+
+        response = await model.generate(request)
+
+        print("✅ Base64 image analysis successful")
+        print(f"   Response: {response.content}")
+
+    except Exception as e:
+        print(f"❌ Base64 image test failed: {e}")
+        raise
+
+
+async def run_async_tests() -> None:
+    """Run async tests."""
+    try:
+        # await test_text_generation()
+        # await test_streaming_generation()
+        # await test_embedding_generation()
+        await test_multimodal_vision()
+        await test_multimodal_with_base64()
+    except Exception as e:
+        print(f"\n❌ Async tests failed: {e}")
+        raise
+
+
+def main() -> None:
+    """Run all tests."""
+    print("\n")
+    print("╔" + "=" * 58 + "╗")
+    print("║" + " " * 58 + "║")
+    print("║" + "  OLLAMA LLM PROVIDER INTEGRATION TESTS".center(58) + "║")
+    print("║" + " " * 58 + "║")
+    print("╚" + "=" * 58 + "╝")
+
+    try:
+        # Sync tests
+        # test_ollama_model_creation()
+        # test_model_configuration()
+        # test_embedding_model_creation()
+        # test_factory_provider_support()
+        # test_openai_model_direct_creation()
+
+        # Async tests (actual API calls)
+        print("\n" + "=" * 60)
+        print("Running async tests (actual API calls)...")
+        print("=" * 60)
+        asyncio.run(run_async_tests())
+
+        print("\n" + "=" * 60)
+        print("✅ ALL TESTS PASSED!")
+        print("=" * 60 + "\n")
+
+    except Exception as e:
+        print("\n" + "=" * 60)
+        print(f"❌ TESTS FAILED: {e}")
+        print("=" * 60 + "\n")
+        raise
+
+
+if __name__ == "__main__":
+    main()