Skip to content

Ollama prompt driver

OllamaPromptDriver

Bases: BasePromptDriver

Attributes:

Name Type Description
model str

Model name.
Source code in griptape/drivers/prompt/ollama_prompt_driver.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
@define
class OllamaPromptDriver(BasePromptDriver):
    """
    Attributes:
        model: Model name.
    """

    model: str = field(kw_only=True, metadata={"serializable": True})
    host: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
    client: Client = field(
        default=Factory(lambda self: import_optional_dependency("ollama").Client(host=self.host), takes_self=True),
        kw_only=True,
    )
    tokenizer: BaseTokenizer = field(
        default=Factory(
            lambda self: SimpleTokenizer(
                characters_per_token=4, max_input_tokens=2000, max_output_tokens=self.max_tokens
            ),
            takes_self=True,
        ),
        kw_only=True,
    )
    options: dict = field(
        default=Factory(
            lambda self: {
                "temperature": self.temperature,
                "stop": self.tokenizer.stop_sequences,
                "num_predict": self.max_tokens,
            },
            takes_self=True,
        ),
        kw_only=True,
    )

    def try_run(self, prompt_stack: PromptStack) -> TextArtifact:
        response = self.client.chat(**self._base_params(prompt_stack))

        if isinstance(response, dict):
            return TextArtifact(value=response["message"]["content"])
        else:
            raise Exception("invalid model response")

    def try_stream(self, prompt_stack: PromptStack) -> Iterator[TextArtifact]:
        stream = self.client.chat(**self._base_params(prompt_stack), stream=True)

        if isinstance(stream, Iterator):
            for chunk in stream:
                yield TextArtifact(value=chunk["message"]["content"])
        else:
            raise Exception("invalid model response")

    def _base_params(self, prompt_stack: PromptStack) -> dict:
        messages = [{"role": input.role, "content": input.content} for input in prompt_stack.inputs]

        return {"messages": messages, "model": self.model, "options": self.options}

client: Client = field(default=Factory(lambda self: import_optional_dependency('ollama').Client(host=self.host), takes_self=True), kw_only=True) class-attribute instance-attribute

host: Optional[str] = field(default=None, kw_only=True, metadata={'serializable': True}) class-attribute instance-attribute

model: str = field(kw_only=True, metadata={'serializable': True}) class-attribute instance-attribute

options: dict = field(default=Factory(lambda self: {'temperature': self.temperature, 'stop': self.tokenizer.stop_sequences, 'num_predict': self.max_tokens}, takes_self=True), kw_only=True) class-attribute instance-attribute

tokenizer: BaseTokenizer = field(default=Factory(lambda self: SimpleTokenizer(characters_per_token=4, max_input_tokens=2000, max_output_tokens=self.max_tokens), takes_self=True), kw_only=True) class-attribute instance-attribute

try_run(prompt_stack)

Source code in griptape/drivers/prompt/ollama_prompt_driver.py 
49
50
51
52
53
54
55
def try_run(self, prompt_stack: PromptStack) -> TextArtifact:
    response = self.client.chat(**self._base_params(prompt_stack))

    if isinstance(response, dict):
        return TextArtifact(value=response["message"]["content"])
    else:
        raise Exception("invalid model response")

try_stream(prompt_stack)

Source code in griptape/drivers/prompt/ollama_prompt_driver.py 
57
58
59
60
61
62
63
64
def try_stream(self, prompt_stack: PromptStack) -> Iterator[TextArtifact]:
    stream = self.client.chat(**self._base_params(prompt_stack), stream=True)

    if isinstance(stream, Iterator):
        for chunk in stream:
            yield TextArtifact(value=chunk["message"]["content"])
    else:
        raise Exception("invalid model response")