[FEATURE] Ollama Provider Finalisation

examples/agents/ollama_agent.py

@@ -0,0 +1,93 @@
+"""Ollama Agent Example
+
+This example demonstrates how to create a simple agent using Legion's decorator syntax.
+The agent can use both internal (nested) tools and external tools.
+"""
+
+from typing import Annotated, List
+
+from dotenv import load_dotenv
+from pydantic import Field
+
+from legion.agents import agent
+from legion.interface.decorators import tool
+
+load_dotenv()
+
+
+@tool
+def add_numbers(
+    numbers: Annotated[List[float], Field(description="List of numbers to add together")]
+) -> float:
+    """Add a list of numbers together and return the sum."""
+    return sum(numbers)
+
+
+@tool
+def multiply(
+    a: Annotated[float, Field(description="First number to multiply")],
+    b: Annotated[float, Field(description="Second number to multiply")]
+) -> float:
+    """Multiply two numbers together."""
+    return a * b
+
+
+@agent(
+    model="ollama:llama3.2",
+    temperature=0.3,
+    tools=[add_numbers, multiply],  # Bind external tools
+)
+class MathHelper:
+    """An agent that helps with basic arithmetic and string operations.
+
+    I can perform calculations and manipulate text based on your requests.
+    I have both external tools for math operations and internal tools for formatting.
+    I format arguments in JSON format and pay attention to decimal places.
+    My answers are succinct and only contain the input arguments and the answer.
+    """
+
+    @tool
+    def format_result(
+        self,
+        number: Annotated[float, Field(description="Number to format")],
+        prefix: Annotated[str, Field(description="Text to add before the number")] = (
+            "Result: "
+        )
+    ) -> str:
+        """Format a number with a custom prefix."""
+        return f"{prefix}{number:.2f}"
+
+
+async def main():
+    # Create an instance of our agent
+    agent = MathHelper()
+
+    # Example 1: Using external add_numbers tool with internal format_result
+    response = await agent.aprocess(
+        "I have the numbers 1.5, 2.5, and 3.5. Can you add them together and format "
+        "the result nicely?"
+    )
+    print("Example 1 Response:")
+    print(response.content)
+    print()
+
+    # Example 2: Using external multiply tool with internal format_result
+    response = await agent.aprocess(
+        "I have the numbers 4.2 and 2.0. Can you multiply them and then format "
+        "the result nicely?"
+    )
+    print("Example 2 Response:")
+    print(response.content)
+    print()
+
+    # Example 3: Complex operation using both external and internal tools
+    response = await agent.aprocess(
+        "I need to add the numbers 10.5 and 20.5, then multiply the result by 2"
+    )
+    print("Example 3 Response:")
+    print(response.content)
+
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(main())

legion/providers/ollama.py

@@ -1,5 +1,6 @@
 """Ollama-specific implementation of the LLM interface"""
 
+import ast
 import json
 from typing import Any, Dict, List, Optional, Sequence, Type
 
@@ -8,7 +9,6 @@
 from ..errors import ProviderError
 from ..interface.base import LLMInterface
 from ..interface.schemas import (
-    ChatParameters,
     Message,
     ModelResponse,
     ProviderConfig,
@@ -29,6 +29,11 @@ def create_provider(self, config: Optional[ProviderConfig] = None, **kwargs) ->
 class OllamaProvider(LLMInterface):
     """Ollama-specific provider implementation"""
 
+    def __init__(self, config: ProviderConfig, debug: bool = False):
+        """Initialize provider with both sync and async clients"""
+        super().__init__(config, debug)
+        self._async_client = None  # Initialize async client lazily
+
     def _setup_client(self) -> None:
         """Initialize Ollama client"""
         try:
@@ -66,24 +71,40 @@ def _format_messages(self, messages: List[Message]) -> List[Dict[str, Any]]:
 
         return ollama_messages
 
+    def _format_arguments(
+            self,
+            arguments: Any
+    ):
+        """Convert Tools Arguments into JSON"""
+        if isinstance(arguments, dict):
+            try:
+                args = {k:ast.literal_eval(v) for k,v in arguments.items()}
+            except Exception:
+                args = arguments
+        else:
+                args = arguments
+        return json.dumps(args)
+
+
     def _get_chat_completion(
         self,
         messages: List[Message],
         model: str,
-        params: ChatParameters
+        temperature: float,
+        stream: Optional[bool] = False
     ) -> ModelResponse:
         """Get a basic chat completion"""
         try:
             # Build options dictionary
             options = {
-                "temperature": params.temperature
+                "temperature": temperature
             }
 
             response = self.client.chat(
                 model=model,
                 messages=self._format_messages(messages),
                 options=options,
-                stream=params.stream
+                stream=stream
             )
 
             return ModelResponse(
@@ -238,3 +259,225 @@ def _extract_tool_calls(self, response: Any) -> Optional[List[Dict[str, Any]]]:
                 print(f"\nExtracted tool call: {json.dumps(tool_calls[-1], indent=2)}")
 
         return tool_calls if tool_calls else None
+
+    async def _asetup_client(self) -> None:
+        """Initialize async Ollama client"""
+        from ollama import AsyncClient
+        try:
+            self._async_client = AsyncClient(
+                host=self.config.base_url or "http://localhost:11434",
+            )
+        except Exception as e:
+            raise ProviderError(f"Failed to initialize async Ollama client: {str(e)}")
+
+    async def _ensure_async_client(self) -> None:
+        """Ensure async client is initialized"""
+        if self._async_client is None:
+            await self._asetup_client()
+
+    async def _aget_chat_completion(self, messages, model, temperature, max_tokens = None):
+        """Get a basic chat completion asynchronously"""
+        try:
+            await self._ensure_async_client()
+            response = await self._async_client.chat(
+                model=model,
+                messages=[msg.model_dump() for msg in messages],
+                options={"temperature": temperature}
+            )
+
+            return ModelResponse(
+                content=self._extract_content(response),
+                raw_response=self._response_to_dict(response),
+                usage=self._extract_usage(response),
+                tool_calls=None
+            )
+        except Exception as e:
+            raise ProviderError(f"Ollama async completion failed: {str(e)}")
+
+    async def _aget_json_completion(
+            self,
+            messages,
+            model,
+            schema,
+            temperature,
+            max_tokens = None,
+            preserve_tool_calls: Optional[List[Dict[str, Any]]] = None
+        ):
+        """Get a chat completion formatted as JSON asynchronously"""
+        try:
+            await self._ensure_async_client()
+
+            # Get generic JSON formatting prompt
+            formatting_prompt = self._get_json_formatting_prompt(schema, messages[-1].content)
+
+            # Create messages for Ollama
+            messages = [
+                {"role": "system", "content": formatting_prompt}
+            ]
+
+            # Add remaining messages, skipping system
+            messages.extend([
+                msg for msg in messages
+                if msg["role"] != Role.SYSTEM
+            ])
+
+            response = await self._async_client.chat(
+                model=model,
+                messages=messages,
+                format="json",
+                options={"temperature": temperature}
+            )
+
+            # Validate against schema
+            content = response.message.content
+            try:
+                data = json.loads(content)
+                schema.model_validate(data)
+            except Exception as e:
+                raise ProviderError(f"Invalid JSON response: {str(e)}")
+
+            return ModelResponse(
+                content=content,
+                raw_response=self._response_to_dict(response),
+                usage=self._extract_usage(response),
+                tool_calls=preserve_tool_calls
+            )
+        except Exception as e:
+            raise ProviderError(f"Ollama async JSON completion failed: {str(e)}")
+
+    async def _aget_tool_completion(
+            self,
+            messages,
+            model,
+            tools,
+            temperature,
+            max_tokens = None,
+            format_json = False,
+            json_schema = None
+        ):
+        "Get completion with tool usage asynchronously"""
+        await self._ensure_async_client()
+        current_messages = list(messages)
+        all_tool_calls = []
+
+        try:
+            # First phase: Use tools
+            while True:
+                if self.debug:
+                    print("\n🔄 Making async Ollama API call:")
+                    print(f"Messages count: {len(current_messages)}")
+                    print("Tools:", [t.name for t in tools])
+
+                # Convert messages and tools to dict format
+                message_dicts = self._format_messages(current_messages)
+                tool_dicts = [t.model_dump() for t in tools]
+
+                try:
+                    response = await self._async_client.chat(
+                        model=model,
+                        messages=message_dicts,
+                        tools=tool_dicts,
+                        options={"temperature": temperature}
+                    )
+                except Exception as api_error:
+                    if self.debug:
+                        print(f"\n❌ Async API call failed: {str(api_error)}")
+                    raise
+
+                content = response.message.content or ""
+
+                # Process tool calls if any
+                if response.message.tool_calls:
+                    # First add the assistant's message with tool calls
+                    tool_call_data = []
+                    for id, tool_call in enumerate(response.message.tool_calls):
+                        call_data = {
+                            "id": str(id),
+                            "type": "function",
+                            "function": {
+                                "name": tool_call.function.name,
+                                "arguments": self._format_arguments(tool_call.function.arguments)
+                            }
+                        }
+                        tool_call_data.append(call_data)
+
+                    # Add the assistant's message with tool calls
+                    current_messages.append(Message(
+                        role=Role.ASSISTANT,
+                        content=content,
+                        tool_calls=tool_call_data
+                    ))
+
+                    # Process each tool call
+                    for id, tool_call in enumerate(response.message.tool_calls):
+                        tool = next(
+                            (t for t in tools if t.name == tool_call.function.name),
+                            None
+                        )
+
+                        if tool:
+                            args = json.loads(self._format_arguments(tool_call.function.arguments))
+                            result = await tool.arun(**args)  # Use async tool call
+
+                            if self.debug:
+                                print(f"Tool {tool.name} returned: {result}")
+
+                            # Add the tool's response
+                            current_messages.append(Message(
+                                role=Role.TOOL,
+                                content=json.dumps(result) if isinstance(result, dict) else str(result),
+                                tool_call_id=str(id),
+                                name=tool_call.function.name
+                            ))
+
+                            # Store tool call for final response
+                            call_data = next(
+                                c for c in tool_call_data
+                                if c["id"] == str(id)
+                            )
+                            call_data["result"] = json.dumps(result) if isinstance(result, dict) else str(result)
+                            all_tool_calls.append(call_data)
+                    continue
+
+                # No more tool calls - get final response
+                if format_json and json_schema:
+                    json_response = await self._aget_json_completion(
+                        messages=current_messages,
+                        model=model,
+                        schema=json_schema,
+                        temperature=0.0,
+                        preserve_tool_calls=all_tool_calls if all_tool_calls else None
+                    )
+                    return json_response
+
+                return ModelResponse(
+                    content=content,
+                    raw_response=self._response_to_dict(response),
+                    tool_calls=all_tool_calls if all_tool_calls else None,
+                    usage=self._extract_usage(response)
+                )
+
+        except Exception as e:
+            raise ProviderError(f"Ollama async tool completion failed: {str(e)}")
+
+    def _response_to_dict(self, response: Any) -> Dict[str, Any]:
+        """Convert Ollama response to dictionary"""
+        return {
+            "message": {
+                "role": response.message.role,
+                "content": response.message.content,
+                "tool_calls": [
+                    {
+                        "id": tool_call.id,
+                        "type": tool_call.type,
+                        "function": {
+                            "name": tool_call.function.name,
+                            "arguments": tool_call.function.arguments
+                        }
+                    }
+                    for tool_call in (response.message.tool_calls or [])
+                ] if response.message.tool_calls else None
+            },
+            "model": response.model,
+            "created_at": response.created_at
+        }