From 30c962ef43344ca28e1068af9cfd1c3c4b84a359 Mon Sep 17 00:00:00 2001
From: Dirk Brand <51947788+dirkbrnd@users.noreply.github.com>
Date: Fri, 20 Dec 2024 16:04:35 +0200
Subject: [PATCH] Add model_copy to make ollama chat deep copy work with a set
 client (#1603)

## Description
If you set the client manually for Ollama, then there is an issue when
the agent model is copied where certain properties can't be pickled.
This resolves it.
---
 cookbook/agents/14_generate_image.py          |  8 ++------
 cookbook/providers/ollama/agent_set_client.py | 18 ++++++++++++++++++
 phi/model/ollama/chat.py                      |  4 ++++
 3 files changed, 24 insertions(+), 6 deletions(-)
 create mode 100644 cookbook/providers/ollama/agent_set_client.py

diff --git a/cookbook/agents/14_generate_image.py b/cookbook/agents/14_generate_image.py
index 1f8323155..642d62710 100644
--- a/cookbook/agents/14_generate_image.py
+++ b/cookbook/agents/14_generate_image.py
@@ -16,9 +16,5 @@
 images = image_agent.get_images()
 if images and isinstance(images, list):
     for image_response in images:
-        image_data = image_response.get("data")  # type: ignore
-        if image_data:
-            for image in image_data:
-                image_url = image.get("url")  # type: ignore
-                if image_url:
-                    print(image_url)
+        image_url = image_response.url
+        print(image_url)
diff --git a/cookbook/providers/ollama/agent_set_client.py b/cookbook/providers/ollama/agent_set_client.py
new file mode 100644
index 000000000..3377811e6
--- /dev/null
+++ b/cookbook/providers/ollama/agent_set_client.py
@@ -0,0 +1,18 @@
+"""Run `pip install yfinance` to install dependencies."""
+
+from ollama import Client as OllamaClient
+from phi.agent import Agent, RunResponse  # noqa
+from phi.model.ollama import Ollama
+from phi.playground import Playground, serve_playground_app
+from phi.tools.yfinance import YFinanceTools
+
+agent = Agent(
+    model=Ollama(id="llama3.1:8b", client=OllamaClient()),
+    tools=[YFinanceTools(stock_price=True)],
+    markdown=True,
+)
+
+app = Playground(agents=[agent]).get_app()
+
+if __name__ == "__main__":
+    serve_playground_app("agent_set_client:app", reload=True)
diff --git a/phi/model/ollama/chat.py b/phi/model/ollama/chat.py
index acc744454..51f39089e 100644
--- a/phi/model/ollama/chat.py
+++ b/phi/model/ollama/chat.py
@@ -722,3 +722,7 @@ async def aresponse_stream(self, messages: List[Message]) -> Any:
             async for post_tool_call_response in self.ahandle_post_tool_call_messages_stream(messages=messages):
                 yield post_tool_call_response
         logger.debug("---------- Ollama Async Response End ----------")
+
+    def model_copy(self, *, update: Optional[dict[str, Any]] = None, deep: bool = False) -> "Ollama":
+        new_model = Ollama(**self.model_dump(exclude={"client"}), client=self.client)
+        return new_model