Merge pull request #37 from patrickfleith/feature/33-implement-ollama-provider

Feature/33 implement ollama provider

Author: patrickfleith

README.md

@@ -16,12 +16,22 @@ Currently we support the following dataset types:
 
 - ✅ Text Classification
 - ✅ Raw Text Generation
-- ✅ Instruction Dataset
-    - ✅  UltraChat method
+- ✅ Instruction Dataset (UltraChat-like)
 - [ ] Preference Dataset
 - 📋 More coming soon!
 
-⭐️ Star me if this is something you like!  
+⭐️ Star me if this is something you like! 🌟
+
+
+## Supported LLM Providers
+
+Currently we support the following LLM providers:
+
+- ✔︎ OpenAI
+- ✔︎ Anthropic
+- ✔︎ Google
+- ✔︎ Ollama
+- ✔︎ HF Endpoints (buggy!)
 
 ## Key Features
 

datafast/examples/test_llm_providers.py

@@ -128,6 +128,9 @@ def main():
 
     # Test OpenAI (GPT-4)
     test_provider("openai", "gpt-4o-mini")
+    
+    # Test Ollama (local LLM)
+    test_provider("ollama", "gemma3:4b")
 
 
 if __name__ == "__main__":

datafast/examples/test_ollama_provider.py

@@ -0,0 +1,55 @@
+from datafast.llms import create_provider
+from pydantic import BaseModel, Field
+from typing import Optional
+import sys
+
+"""
+A simple test script for the OllamaProvider.
+This script requires Ollama to be installed and running locally.
+
+You can install Ollama from https://ollama.com/ and then run:
+ollama pull gemma3:4b
+"""
+
+
+class SimpleResponse(BaseModel):
+    """A simple response model with minimal fields to test the OllamaProvider."""
+    answer: str = Field(..., description="The answer to the question")
+    reasoning: str = Field(..., description="The reasoning behind the answer")
+
+
+def test_ollama(model_id: str = "gemma3:4b"):
+    """Test the OllamaProvider with a simple query."""
+    print(f"\n{'=' * 50}")
+    print(f"Testing Ollama provider with model {model_id}")
+    print("=" * 50)
+
+    try:
+        # Create the provider
+        provider = create_provider("ollama", model_id)
+        
+        # Simple test prompt
+        prompt = "What is the capital of France? Provide a short answer and brief reasoning."
+        
+        print(f"Sending prompt: {prompt}")
+        print("Waiting for response (this might take a bit)...")
+        
+        # Generate response
+        response = provider.generate(prompt, SimpleResponse)
+        
+        print("\nResponse received:")
+        print(f"Answer: {response.answer}")
+        print(f"Reasoning: {response.reasoning}")
+        print("\nTest successful!")
+        
+    except Exception as e:
+        print(f"Error testing Ollama provider: {str(e)}")
+        return False
+    
+    return True
+
+
+if __name__ == "__main__":
+    # Get model from command line arguments if provided
+    model = sys.argv[1] if len(sys.argv) > 1 else "gemma3:4b"
+    test_ollama(model)

datafast/llms.py

@@ -203,13 +203,58 @@ def _generate_impl(self, prompt: str | list[dict[str, str]], response_format: ty
         )
 
 
+class OllamaProvider(LLMProvider):
+    """Ollama provider for structured text generation."""
+
+    # No API key needed for local Ollama
+    DEFAULT_MODEL = "llama3:latest"
+
+    @property
+    def name(self) -> str:
+        return "ollama"
+        
+    def _get_api_key(self) -> str:
+        """Override _get_api_key since Ollama doesn't need an API key"""
+        return "not_needed"  # Return a dummy value
+
+    def _initialize_client(self):
+        try:
+            import ollama
+            return ollama
+        except ImportError as e:
+            raise ImportError(f"Ollama package not installed. Install it with 'pip install ollama': {str(e)}")
+        except Exception as e:
+            raise ValueError(f"Error initializing Ollama client: {str(e)}")
+
+    def _generate_impl(
+        self, prompt: str | list[dict[str, str]], response_format: type[BaseModel]
+    ) -> BaseModel:
+        # Convert prompt to messages format if it's a string
+        messages = get_messages(prompt) if isinstance(prompt, str) else prompt
+        
+        # Get schema for the response format
+        schema = response_format.model_json_schema()
+        
+        # Call the Ollama chat API
+        response = self.client.chat(
+            messages=messages,
+            model=self.model_id,
+            format=schema,
+        )
+        
+        # Parse the response content and validate against the Pydantic model
+        # Unlike other providers that use instructor and return the parsed model directly,
+        # we need to manually parse the JSON response here
+        return response_format.model_validate_json(response.message.content)
+
+
 def create_provider(
     provider: str, model_id: str | None = None, **kwargs
 ) -> LLMProvider:
     """Create an LLM provider for structured text generation.
 
     Args:
-        provider: Provider name ('anthropic', 'google', or 'openai')
+        provider: Provider name ('anthropic', 'google', 'openai', 'ollama')
         model_id: Optional model identifier. If not provided, uses provider's default
         **kwargs: Additional provider-specific arguments
 
@@ -220,6 +265,7 @@ def create_provider(
         "anthropic": AnthropicProvider,
         "google": GoogleProvider,
         "openai": OpenAIProvider,
+        "ollama": OllamaProvider,
     }
 
     provider_class = provider_map.get(provider.lower())