From 40cdf58a71afa8270783e50eb83cb82f7509f2ef Mon Sep 17 00:00:00 2001
From: Arthur Brenno <64020210+arthurbrenno@users.noreply.github.com>
Date: Sat, 8 Feb 2025 10:05:50 -0300
Subject: [PATCH] .

---
 pyproject.toml                                |  14 ++-
 .../llms/integrations/ollama/__init__.py      |   0
 .../llms/integrations/ollama/ollama.py        | 105 ++++++++++++++++++
 src/intellibricks/llms/types.py               |   7 +-
 uv.lock                                       |  17 ++-
 5 files changed, 140 insertions(+), 3 deletions(-)
 create mode 100644 src/intellibricks/llms/integrations/ollama/__init__.py
 create mode 100644 src/intellibricks/llms/integrations/ollama/ollama.py

diff --git a/pyproject.toml b/pyproject.toml
index 4576fb3..ea3eb34 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,19 @@ version = "v0.8.3"
 description = "IntelliBricks provides a streamlined set of tools for developing AI-powered applications. It simplifies complex tasks such as interacting with LLMs, training machine learning models, and implementing Retrieval Augmented Generation (RAG). Focus on building your application logic, not wrestling with boilerplate. IntelliBricks empowers you to build intelligent applications faster and more efficiently."
 readme = "README.md"
 requires-python = ">=3.13.0"
-dependencies = ["langfuse>=2.57.0", "groq>=0.13.1", "google-genai>=0.3.0", "msgspec", "langchain-core==0.3.28", "openai>=1.54.3", "beautifulsoup4>=4.12.3", "mypy>=1.14.1", "python-dotenv>=1.0.1", "architecture>=0.5.23"]
+dependencies = [
+    "langfuse>=2.57.0",
+ "groq>=0.13.1",
+ "google-genai>=0.3.0",
+ "msgspec",
+ "langchain-core==0.3.28",
+ "openai>=1.54.3",
+ "beautifulsoup4>=4.12.3",
+ "mypy>=1.14.1",
+ "python-dotenv>=1.0.1",
+ "architecture>=0.5.23",
+ "ollama>=0.4.7",
+]
 
 [project.urls]
 Source = "https://github.com/arthurbrenno/intellibricks"
diff --git a/src/intellibricks/llms/integrations/ollama/__init__.py b/src/intellibricks/llms/integrations/ollama/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/intellibricks/llms/integrations/ollama/ollama.py b/src/intellibricks/llms/integrations/ollama/ollama.py
new file mode 100644
index 0000000..0a1b231
--- /dev/null
+++ b/src/intellibricks/llms/integrations/ollama/ollama.py
@@ -0,0 +1,105 @@
+# import timeit
+# from typing import Literal, Optional, Sequence, TypeVar, overload, override
+
+# import msgspec
+# from architecture.utils.decorators import ensure_module_installed
+# from langfuse.client import os
+# from intellibricks.llms.util import ms_type_to_schema
+# from intellibricks.llms.base import (
+#     LanguageModel,
+# )
+# from intellibricks.llms.constants import FinishReason
+# from intellibricks.llms.types import (
+#     AudioTranscription,
+#     CalledFunction,
+#     ChatCompletion,
+#     CompletionTokensDetails,
+#     Function,
+#     GeneratedAssistantMessage,
+#     Message,
+#     MessageChoice,
+#     OpenAIModelType,
+#     Part,
+#     PromptTokensDetails,
+#     RawResponse,
+#     SentenceSegment,
+#     ToolCall,
+#     ToolCallSequence,
+#     ToolInputType,
+#     TypeAlias,
+#     Usage,
+# )
+# from ollama import AsyncClient
+
+# S = TypeVar("S", bound=msgspec.Struct, default=RawResponse)
+# DeepSeekModels = Literal[
+#     "deepseek-r1:1.5b",
+#     "deepseek-r1:7b",
+#     "deepseek-r1:8b",
+#     "deepseek-r1:14b",
+#     "deepseek-r1:32b",
+# ]
+
+# ChatModel = Literal[DeepSeekModels]
+
+
+# class OllamaLanguageModel(LanguageModel, frozen=True):
+#     model_name: ChatModel
+#     max_retries: int = 2
+
+#     @overload
+#     async def chat_async(
+#         self,
+#         messages: Sequence[Message],
+#         *,
+#         response_model: None = None,
+#         n: Optional[int] = None,
+#         temperature: Optional[float] = None,
+#         max_completion_tokens: Optional[int] = None,
+#         top_p: Optional[float] = None,
+#         top_k: Optional[int] = None,
+#         stop_sequences: Optional[Sequence[str]] = None,
+#         tools: Optional[Sequence[ToolInputType]] = None,
+#         timeout: Optional[float] = None,
+#     ) -> ChatCompletion[RawResponse]: ...
+#     @overload
+#     async def chat_async(
+#         self,
+#         messages: Sequence[Message],
+#         *,
+#         response_model: type[S],
+#         n: Optional[int] = None,
+#         temperature: Optional[float] = None,
+#         max_completion_tokens: Optional[int] = None,
+#         top_p: Optional[float] = None,
+#         top_k: Optional[int] = None,
+#         stop_sequences: Optional[Sequence[str]] = None,
+#         tools: Optional[Sequence[ToolInputType]] = None,
+#         timeout: Optional[float] = None,
+#     ) -> ChatCompletion[S]: ...
+
+#     @ensure_module_installed("ollama", "ollama")
+#     @override
+#     async def chat_async(
+#         self,
+#         messages: Sequence[Message],
+#         *,
+#         response_model: Optional[type[S]] = None,
+#         n: Optional[int] = None,
+#         temperature: Optional[float] = None,
+#         max_completion_tokens: Optional[int] = None,
+#         top_p: Optional[float] = None,
+#         top_k: Optional[int] = None,
+#         stop_sequences: Optional[Sequence[str]] = None,
+#         tools: Optional[Sequence[ToolInputType]] = None,
+#         timeout: Optional[float] = None,
+#     ) -> ChatCompletion[S] | ChatCompletion[RawResponse]:
+#         now = timeit.default_timer()
+#         client = AsyncClient()
+#         completion = await client.chat(
+#             model=self.model_name,
+#             messages=[m.to_ollama_message() for m in messages],
+#             format=ms_type_to_schema(response_model, openai_like=True)
+#             if response_model is not None
+#             else None,
+#         )
diff --git a/src/intellibricks/llms/types.py b/src/intellibricks/llms/types.py
index 99b60f9..35e02df 100644
--- a/src/intellibricks/llms/types.py
+++ b/src/intellibricks/llms/types.py
@@ -111,7 +111,7 @@
 
 import msgspec
 from architecture import dp, log
-from architecture.data.files import find_extension, ext_to_mime, bytes_to_mime
+from architecture.data.files import bytes_to_mime, ext_to_mime, find_extension
 from architecture.utils.decorators import ensure_module_installed
 
 from intellibricks.llms.util import (
@@ -155,6 +155,7 @@
     from groq.types.shared_params.function_definition import (
         FunctionDefinition as GroqFunctionDefinition,
     )
+    from ollama._types import Message as OllamaMessage
     from openai.types.chat.chat_completion_content_part_param import (
         ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam,
     )
@@ -1897,6 +1898,10 @@ def to_google_format(self) -> GenaiContent: ...
     def to_openai_format(self) -> ChatCompletionMessageParam:
         raise NotImplementedError
 
+    # @abstractmethod
+    def to_ollama_message(self) -> OllamaMessage:
+        raise NotImplementedError
+
     # @abstractmethod
     def to_groq_format(self) -> GroqChatCompletionMessageParam:
         raise NotImplementedError
diff --git a/uv.lock b/uv.lock
index f4567f4..4c2d1d2 100644
--- a/uv.lock
+++ b/uv.lock
@@ -871,7 +871,7 @@ wheels = [
 
 [[package]]
 name = "intellibricks"
-version = "0.7.25"
+version = "0.8.3"
 source = { editable = "." }
 dependencies = [
     { name = "architecture" },
@@ -882,6 +882,7 @@ dependencies = [
     { name = "langfuse" },
     { name = "msgspec" },
     { name = "mypy" },
+    { name = "ollama" },
     { name = "openai" },
     { name = "python-dotenv" },
 ]
@@ -933,6 +934,7 @@ requires-dist = [
     { name = "langfuse", specifier = ">=2.57.0" },
     { name = "msgspec", git = "https://github.com/jcrist/msgspec.git?rev=main" },
     { name = "mypy", specifier = ">=1.14.1" },
+    { name = "ollama", specifier = ">=0.4.7" },
     { name = "openai", specifier = ">=1.54.3" },
     { name = "python-dotenv", specifier = ">=1.0.1" },
 ]
@@ -1755,6 +1757,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7b/9c/4fce9cf39dde2562584e4cfd351a0140240f82c0e3569ce25a250f47037d/numpy-2.2.1-cp313-cp313t-win_amd64.whl", hash = "sha256:bff7d8ec20f5f42607599f9994770fa65d76edca264a87b5e4ea5629bce12268", size = 12693107 },
 ]
 
+[[package]]
+name = "ollama"
+version = "0.4.7"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "httpx" },
+    { name = "pydantic" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b0/6d/dc77539c735bbed5d0c873fb029fb86aa9f0163df169b34152914331c369/ollama-0.4.7.tar.gz", hash = "sha256:891dcbe54f55397d82d289c459de0ea897e103b86a3f1fad0fdb1895922a75ff", size = 12843 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/31/83/c3ffac86906c10184c88c2e916460806b072a2cfe34cdcaf3a0c0e836d39/ollama-0.4.7-py3-none-any.whl", hash = "sha256:85505663cca67a83707be5fb3aeff0ea72e67846cea5985529d8eca4366564a1", size = 13210 },
+]
+
 [[package]]
 name = "openai"
 version = "1.60.0"