initial commit

venv/Lib/site-packages/ollama/__init__.py

@@ -0,0 +1,59 @@
+from ollama._client import AsyncClient, Client
+from ollama._types import (
+  ChatResponse,
+  EmbeddingsResponse,
+  EmbedResponse,
+  GenerateResponse,
+  Image,
+  ListResponse,
+  Message,
+  Options,
+  ProcessResponse,
+  ProgressResponse,
+  RequestError,
+  ResponseError,
+  ShowResponse,
+  StatusResponse,
+  Tool,
+  WebFetchResponse,
+  WebSearchResponse,
+)
+
+__all__ = [
+  'AsyncClient',
+  'ChatResponse',
+  'Client',
+  'EmbedResponse',
+  'EmbeddingsResponse',
+  'GenerateResponse',
+  'Image',
+  'ListResponse',
+  'Message',
+  'Options',
+  'ProcessResponse',
+  'ProgressResponse',
+  'RequestError',
+  'ResponseError',
+  'ShowResponse',
+  'StatusResponse',
+  'Tool',
+  'WebFetchResponse',
+  'WebSearchResponse',
+]
+
+_client = Client()
+
+generate = _client.generate
+chat = _client.chat
+embed = _client.embed
+embeddings = _client.embeddings
+pull = _client.pull
+push = _client.push
+create = _client.create
+delete = _client.delete
+list = _client.list
+copy = _client.copy
+show = _client.show
+ps = _client.ps
+web_search = _client.web_search
+web_fetch = _client.web_fetch

venv/Lib/site-packages/ollama/_types.py

@@ -0,0 +1,630 @@
+import contextlib
+import json
+from base64 import b64decode, b64encode
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Union
+
+from pydantic import (
+  BaseModel,
+  ByteSize,
+  ConfigDict,
+  Field,
+  model_serializer,
+)
+from pydantic.json_schema import JsonSchemaValue
+from typing_extensions import Annotated, Literal
+
+
+class SubscriptableBaseModel(BaseModel):
+  def __getitem__(self, key: str) -> Any:
+    """
+    >>> msg = Message(role='user')
+    >>> msg['role']
+    'user'
+    >>> msg = Message(role='user')
+    >>> msg['nonexistent']
+    Traceback (most recent call last):
+    KeyError: 'nonexistent'
+    """
+    if key in self:
+      return getattr(self, key)
+
+    raise KeyError(key)
+
+  def __setitem__(self, key: str, value: Any) -> None:
+    """
+    >>> msg = Message(role='user')
+    >>> msg['role'] = 'assistant'
+    >>> msg['role']
+    'assistant'
+    >>> tool_call = Message.ToolCall(function=Message.ToolCall.Function(name='foo', arguments={}))
+    >>> msg = Message(role='user', content='hello')
+    >>> msg['tool_calls'] = [tool_call]
+    >>> msg['tool_calls'][0]['function']['name']
+    'foo'
+    """
+    setattr(self, key, value)
+
+  def __contains__(self, key: str) -> bool:
+    """
+    >>> msg = Message(role='user')
+    >>> 'nonexistent' in msg
+    False
+    >>> 'role' in msg
+    True
+    >>> 'content' in msg
+    False
+    >>> msg.content = 'hello!'
+    >>> 'content' in msg
+    True
+    >>> msg = Message(role='user', content='hello!')
+    >>> 'content' in msg
+    True
+    >>> 'tool_calls' in msg
+    False
+    >>> msg['tool_calls'] = []
+    >>> 'tool_calls' in msg
+    True
+    >>> msg['tool_calls'] = [Message.ToolCall(function=Message.ToolCall.Function(name='foo', arguments={}))]
+    >>> 'tool_calls' in msg
+    True
+    >>> msg['tool_calls'] = None
+    >>> 'tool_calls' in msg
+    True
+    >>> tool = Tool()
+    >>> 'type' in tool
+    True
+    """
+    if key in self.model_fields_set:
+      return True
+
+    if value := self.__class__.model_fields.get(key):
+      return value.default is not None
+
+    return False
+
+  def get(self, key: str, default: Any = None) -> Any:
+    """
+    >>> msg = Message(role='user')
+    >>> msg.get('role')
+    'user'
+    >>> msg = Message(role='user')
+    >>> msg.get('nonexistent')
+    >>> msg = Message(role='user')
+    >>> msg.get('nonexistent', 'default')
+    'default'
+    >>> msg = Message(role='user', tool_calls=[ Message.ToolCall(function=Message.ToolCall.Function(name='foo', arguments={}))])
+    >>> msg.get('tool_calls')[0]['function']['name']
+    'foo'
+    """
+    return getattr(self, key) if hasattr(self, key) else default
+
+
+class Options(SubscriptableBaseModel):
+  # load time options
+  numa: Optional[bool] = None
+  num_ctx: Optional[int] = None
+  num_batch: Optional[int] = None
+  num_gpu: Optional[int] = None
+  main_gpu: Optional[int] = None
+  low_vram: Optional[bool] = None
+  f16_kv: Optional[bool] = None
+  logits_all: Optional[bool] = None
+  vocab_only: Optional[bool] = None
+  use_mmap: Optional[bool] = None
+  use_mlock: Optional[bool] = None
+  embedding_only: Optional[bool] = None
+  num_thread: Optional[int] = None
+
+  # runtime options
+  num_keep: Optional[int] = None
+  seed: Optional[int] = None
+  num_predict: Optional[int] = None
+  top_k: Optional[int] = None
+  top_p: Optional[float] = None
+  tfs_z: Optional[float] = None
+  typical_p: Optional[float] = None
+  repeat_last_n: Optional[int] = None
+  temperature: Optional[float] = None
+  repeat_penalty: Optional[float] = None
+  presence_penalty: Optional[float] = None
+  frequency_penalty: Optional[float] = None
+  mirostat: Optional[int] = None
+  mirostat_tau: Optional[float] = None
+  mirostat_eta: Optional[float] = None
+  penalize_newline: Optional[bool] = None
+  stop: Optional[Sequence[str]] = None
+
+
+class BaseRequest(SubscriptableBaseModel):
+  model: Annotated[str, Field(min_length=1)]
+  'Model to use for the request.'
+
+
+class BaseStreamableRequest(BaseRequest):
+  stream: Optional[bool] = None
+  'Stream response.'
+
+
+class BaseGenerateRequest(BaseStreamableRequest):
+  options: Optional[Union[Mapping[str, Any], Options]] = None
+  'Options to use for the request.'
+
+  format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None
+  'Format of the response.'
+
+  keep_alive: Optional[Union[float, str]] = None
+  'Keep model alive for the specified duration.'
+
+
+class Image(BaseModel):
+  value: Union[str, bytes, Path]
+
+  @model_serializer
+  def serialize_model(self):
+    if isinstance(self.value, (Path, bytes)):
+      return b64encode(self.value.read_bytes() if isinstance(self.value, Path) else self.value).decode()
+
+    if isinstance(self.value, str):
+      try:
+        if Path(self.value).exists():
+          return b64encode(Path(self.value).read_bytes()).decode()
+      except Exception:
+        # Long base64 string can't be wrapped in Path, so try to treat as base64 string
+        pass
+
+      # String might be a file path, but might not exist
+      if self.value.split('.')[-1] in ('png', 'jpg', 'jpeg', 'webp'):
+        raise ValueError(f'File {self.value} does not exist')
+
+      try:
+        # Try to decode to check if it's already base64
+        b64decode(self.value)
+        return self.value
+      except Exception:
+        raise ValueError('Invalid image data, expected base64 string or path to image file') from Exception
+
+
+class GenerateRequest(BaseGenerateRequest):
+  prompt: Optional[str] = None
+  'Prompt to generate response from.'
+
+  suffix: Optional[str] = None
+  'Suffix to append to the response.'
+
+  system: Optional[str] = None
+  'System prompt to prepend to the prompt.'
+
+  template: Optional[str] = None
+  'Template to use for the response.'
+
+  context: Optional[Sequence[int]] = None
+  'Tokenized history to use for the response.'
+
+  raw: Optional[bool] = None
+
+  images: Optional[Sequence[Image]] = None
+  'Image data for multimodal models.'
+
+  think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None
+  'Enable thinking mode (for thinking models).'
+
+  logprobs: Optional[bool] = None
+  'Return log probabilities for generated tokens.'
+
+  top_logprobs: Optional[int] = None
+  'Number of alternative tokens and log probabilities to include per position (0-20).'
+
+
+class BaseGenerateResponse(SubscriptableBaseModel):
+  model: Optional[str] = None
+  'Model used to generate response.'
+
+  created_at: Optional[str] = None
+  'Time when the request was created.'
+
+  done: Optional[bool] = None
+  'True if response is complete, otherwise False. Useful for streaming to detect the final response.'
+
+  done_reason: Optional[str] = None
+  'Reason for completion. Only present when done is True.'
+
+  total_duration: Optional[int] = None
+  'Total duration in nanoseconds.'
+
+  load_duration: Optional[int] = None
+  'Load duration in nanoseconds.'
+
+  prompt_eval_count: Optional[int] = None
+  'Number of tokens evaluated in the prompt.'
+
+  prompt_eval_duration: Optional[int] = None
+  'Duration of evaluating the prompt in nanoseconds.'
+
+  eval_count: Optional[int] = None
+  'Number of tokens evaluated in inference.'
+
+  eval_duration: Optional[int] = None
+  'Duration of evaluating inference in nanoseconds.'
+
+
+class TokenLogprob(SubscriptableBaseModel):
+  token: str
+  'Token text.'
+
+  logprob: float
+  'Log probability for the token.'
+
+
+class Logprob(TokenLogprob):
+  top_logprobs: Optional[Sequence[TokenLogprob]] = None
+  'Most likely tokens and their log probabilities.'
+
+
+class GenerateResponse(BaseGenerateResponse):
+  """
+  Response returned by generate requests.
+  """
+
+  response: str
+  'Response content. When streaming, this contains a fragment of the response.'
+
+  thinking: Optional[str] = None
+  'Thinking content. Only present when thinking is enabled.'
+
+  context: Optional[Sequence[int]] = None
+  'Tokenized history up to the point of the response.'
+
+  logprobs: Optional[Sequence[Logprob]] = None
+  'Log probabilities for generated tokens.'
+
+
+class Message(SubscriptableBaseModel):
+  """
+  Chat message.
+  """
+
+  role: str
+  "Assumed role of the message. Response messages has role 'assistant' or 'tool'."
+
+  content: Optional[str] = None
+  'Content of the message. Response messages contains message fragments when streaming.'
+
+  thinking: Optional[str] = None
+  'Thinking content. Only present when thinking is enabled.'
+
+  images: Optional[Sequence[Image]] = None
+  """
+  Optional list of image data for multimodal models.
+
+  Valid input types are:
+
+  - `str` or path-like object: path to image file
+  - `bytes` or bytes-like object: raw image data
+
+  Valid image formats depend on the model. See the model card for more information.
+  """
+
+  tool_name: Optional[str] = None
+  'Name of the executed tool.'
+
+  class ToolCall(SubscriptableBaseModel):
+    """
+    Model tool calls.
+    """
+
+    class Function(SubscriptableBaseModel):
+      """
+      Tool call function.
+      """
+
+      name: str
+      'Name of the function.'
+
+      arguments: Mapping[str, Any]
+      'Arguments of the function.'
+
+    function: Function
+    'Function to be called.'
+
+  tool_calls: Optional[Sequence[ToolCall]] = None
+  """
+  Tools calls to be made by the model.
+  """
+
+
+class Tool(SubscriptableBaseModel):
+  type: Optional[str] = 'function'
+
+  class Function(SubscriptableBaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    class Parameters(SubscriptableBaseModel):
+      model_config = ConfigDict(populate_by_name=True)
+      type: Optional[Literal['object']] = 'object'
+      defs: Optional[Any] = Field(None, alias='$defs')
+      items: Optional[Any] = None
+      required: Optional[Sequence[str]] = None
+
+      class Property(SubscriptableBaseModel):
+        model_config = ConfigDict(arbitrary_types_allowed=True)
+
+        type: Optional[Union[str, Sequence[str]]] = None
+        items: Optional[Any] = None
+        description: Optional[str] = None
+        enum: Optional[Sequence[Any]] = None
+
+      properties: Optional[Mapping[str, Property]] = None
+
+    parameters: Optional[Parameters] = None
+
+  function: Optional[Function] = None
+
+
+class ChatRequest(BaseGenerateRequest):
+  @model_serializer(mode='wrap')
+  def serialize_model(self, nxt):
+    output = nxt(self)
+    if output.get('tools'):
+      for tool in output['tools']:
+        if 'function' in tool and 'parameters' in tool['function'] and 'defs' in tool['function']['parameters']:
+          tool['function']['parameters']['$defs'] = tool['function']['parameters'].pop('defs')
+    return output
+
+  messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None
+  'Messages to chat with.'
+
+  tools: Optional[Sequence[Tool]] = None
+  'Tools to use for the chat.'
+
+  think: Optional[Union[bool, Literal['low', 'medium', 'high']]] = None
+  'Enable thinking mode (for thinking models).'
+
+  logprobs: Optional[bool] = None
+  'Return log probabilities for generated tokens.'
+
+  top_logprobs: Optional[int] = None
+  'Number of alternative tokens and log probabilities to include per position (0-20).'
+
+
+class ChatResponse(BaseGenerateResponse):
+  """
+  Response returned by chat requests.
+  """
+
+  message: Message
+  'Response message.'
+
+  logprobs: Optional[Sequence[Logprob]] = None
+  'Log probabilities for generated tokens if requested.'
+
+
+class EmbedRequest(BaseRequest):
+  input: Union[str, Sequence[str]]
+  'Input text to embed.'
+
+  truncate: Optional[bool] = None
+  'Truncate the input to the maximum token length.'
+
+  options: Optional[Union[Mapping[str, Any], Options]] = None
+  'Options to use for the request.'
+
+  keep_alive: Optional[Union[float, str]] = None
+
+  dimensions: Optional[int] = None
+  'Dimensions truncates the output embedding to the specified dimension.'
+
+
+class EmbedResponse(BaseGenerateResponse):
+  """
+  Response returned by embed requests.
+  """
+
+  embeddings: Sequence[Sequence[float]]
+  'Embeddings of the inputs.'
+
+
+class EmbeddingsRequest(BaseRequest):
+  prompt: Optional[str] = None
+  'Prompt to generate embeddings from.'
+
+  options: Optional[Union[Mapping[str, Any], Options]] = None
+  'Options to use for the request.'
+
+  keep_alive: Optional[Union[float, str]] = None
+
+
+class EmbeddingsResponse(SubscriptableBaseModel):
+  """
+  Response returned by embeddings requests.
+  """
+
+  embedding: Sequence[float]
+  'Embedding of the prompt.'
+
+
+class PullRequest(BaseStreamableRequest):
+  """
+  Request to pull the model.
+  """
+
+  insecure: Optional[bool] = None
+  'Allow insecure (HTTP) connections.'
+
+
+class PushRequest(BaseStreamableRequest):
+  """
+  Request to pull the model.
+  """
+
+  insecure: Optional[bool] = None
+  'Allow insecure (HTTP) connections.'
+
+
+class CreateRequest(BaseStreamableRequest):
+  @model_serializer(mode='wrap')
+  def serialize_model(self, nxt):
+    output = nxt(self)
+    if 'from_' in output:
+      output['from'] = output.pop('from_')
+    return output
+
+  """
+  Request to create a new model.
+  """
+  quantize: Optional[str] = None
+  from_: Optional[str] = None
+  files: Optional[Dict[str, str]] = None
+  adapters: Optional[Dict[str, str]] = None
+  template: Optional[str] = None
+  license: Optional[Union[str, List[str]]] = None
+  system: Optional[str] = None
+  parameters: Optional[Union[Mapping[str, Any], Options]] = None
+  messages: Optional[Sequence[Union[Mapping[str, Any], Message]]] = None
+
+
+class ModelDetails(SubscriptableBaseModel):
+  parent_model: Optional[str] = None
+  format: Optional[str] = None
+  family: Optional[str] = None
+  families: Optional[Sequence[str]] = None
+  parameter_size: Optional[str] = None
+  quantization_level: Optional[str] = None
+
+
+class ListResponse(SubscriptableBaseModel):
+  class Model(SubscriptableBaseModel):
+    model: Optional[str] = None
+    modified_at: Optional[datetime] = None
+    digest: Optional[str] = None
+    size: Optional[ByteSize] = None
+    details: Optional[ModelDetails] = None
+
+  models: Sequence[Model]
+  'List of models.'
+
+
+class DeleteRequest(BaseRequest):
+  """
+  Request to delete a model.
+  """
+
+
+class CopyRequest(BaseModel):
+  """
+  Request to copy a model.
+  """
+
+  source: str
+  'Source model to copy.'
+
+  destination: str
+  'Destination model to copy to.'
+
+
+class StatusResponse(SubscriptableBaseModel):
+  status: Optional[str] = None
+
+
+class ProgressResponse(StatusResponse):
+  completed: Optional[int] = None
+  total: Optional[int] = None
+  digest: Optional[str] = None
+
+
+class ShowRequest(BaseRequest):
+  """
+  Request to show model information.
+  """
+
+
+class ShowResponse(SubscriptableBaseModel):
+  modified_at: Optional[datetime] = None
+
+  template: Optional[str] = None
+
+  modelfile: Optional[str] = None
+
+  license: Optional[str] = None
+
+  details: Optional[ModelDetails] = None
+
+  modelinfo: Optional[Mapping[str, Any]] = Field(alias='model_info')
+
+  parameters: Optional[str] = None
+
+  capabilities: Optional[List[str]] = None
+
+
+class ProcessResponse(SubscriptableBaseModel):
+  class Model(SubscriptableBaseModel):
+    model: Optional[str] = None
+    name: Optional[str] = None
+    digest: Optional[str] = None
+    expires_at: Optional[datetime] = None
+    size: Optional[ByteSize] = None
+    size_vram: Optional[ByteSize] = None
+    details: Optional[ModelDetails] = None
+    context_length: Optional[int] = None
+
+  models: Sequence[Model]
+
+
+class WebSearchRequest(SubscriptableBaseModel):
+  query: str
+  max_results: Optional[int] = None
+
+
+class WebSearchResult(SubscriptableBaseModel):
+  content: Optional[str] = None
+  title: Optional[str] = None
+  url: Optional[str] = None
+
+
+class WebFetchRequest(SubscriptableBaseModel):
+  url: str
+
+
+class WebSearchResponse(SubscriptableBaseModel):
+  results: Sequence[WebSearchResult]
+
+
+class WebFetchResponse(SubscriptableBaseModel):
+  title: Optional[str] = None
+  content: Optional[str] = None
+  links: Optional[Sequence[str]] = None
+
+
+class RequestError(Exception):
+  """
+  Common class for request errors.
+  """
+
+  def __init__(self, error: str):
+    super().__init__(error)
+    self.error = error
+    'Reason for the error.'
+
+
+class ResponseError(Exception):
+  """
+  Common class for response errors.
+  """
+
+  def __init__(self, error: str, status_code: int = -1):
+    # try to parse content as JSON and extract 'error'
+    # fallback to raw content if JSON parsing fails
+    with contextlib.suppress(json.JSONDecodeError):
+      error = json.loads(error).get('error', error)
+
+    super().__init__(error)
+    self.error = error
+    'Reason for the error.'
+
+    self.status_code = status_code
+    'HTTP status code of the response.'
+
+  def __str__(self) -> str:
+    return f'{self.error} (status code: {self.status_code})'

venv/Lib/site-packages/ollama/_utils.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import inspect
+import re
+from collections import defaultdict
+from typing import Callable, Union
+
+import pydantic
+
+from ollama._types import Tool
+
+
+def _parse_docstring(doc_string: Union[str, None]) -> dict[str, str]:
+  parsed_docstring = defaultdict(str)
+  if not doc_string:
+    return parsed_docstring
+
+  key = str(hash(doc_string))
+  for line in doc_string.splitlines():
+    lowered_line = line.lower().strip()
+    if lowered_line.startswith('args:'):
+      key = 'args'
+    elif lowered_line.startswith(('returns:', 'yields:', 'raises:')):
+      key = '_'
+
+    else:
+      # maybe change to a list and join later
+      parsed_docstring[key] += f'{line.strip()}\n'
+
+  last_key = None
+  for line in parsed_docstring['args'].splitlines():
+    line = line.strip()
+    if ':' in line:
+      # Split the line on either:
+      # 1. A parenthetical expression like (integer) - captured in group 1
+      # 2. A colon :
+      # Followed by optional whitespace. Only split on first occurrence.
+      parts = re.split(r'(?:\(([^)]*)\)|:)\s*', line, maxsplit=1)
+
+      arg_name = parts[0].strip()
+      last_key = arg_name
+
+      # Get the description - will be in parts[1] if parenthetical or parts[-1] if after colon
+      arg_description = parts[-1].strip()
+      if len(parts) > 2 and parts[1]:  # Has parenthetical content
+        arg_description = parts[-1].split(':', 1)[-1].strip()
+
+      parsed_docstring[last_key] = arg_description
+
+    elif last_key and line:
+      parsed_docstring[last_key] += ' ' + line
+
+  return parsed_docstring
+
+
+def convert_function_to_tool(func: Callable) -> Tool:
+  doc_string_hash = str(hash(inspect.getdoc(func)))
+  parsed_docstring = _parse_docstring(inspect.getdoc(func))
+  schema = type(
+    func.__name__,
+    (pydantic.BaseModel,),
+    {
+      '__annotations__': {k: v.annotation if v.annotation != inspect._empty else str for k, v in inspect.signature(func).parameters.items()},
+      '__signature__': inspect.signature(func),
+      '__doc__': parsed_docstring[doc_string_hash],
+    },
+  ).model_json_schema()
+
+  for k, v in schema.get('properties', {}).items():
+    # If type is missing, the default is string
+    types = {t.get('type', 'string') for t in v.get('anyOf')} if 'anyOf' in v else {v.get('type', 'string')}
+    if 'null' in types:
+      schema['required'].remove(k)
+      types.discard('null')
+
+    schema['properties'][k] = {
+      'description': parsed_docstring[k],
+      'type': ', '.join(types),
+    }
+
+  tool = Tool(
+    type='function',
+    function=Tool.Function(
+      name=func.__name__,
+      description=schema.get('description', ''),
+      parameters=Tool.Function.Parameters(**schema),
+    ),
+  )
+
+  return Tool.model_validate(tool)