[Feature] Add function calling

pyproject.toml

@@ -31,5 +31,6 @@ requires-python = ">=3.10"
 dependencies = [
   "asyncio>=3.4.3",
   "mcp[cli]>=1.9.4",
+  "ollama>=0.5.1",
   "pydantic>=2.11.7"
 ]

src/llm_agents_from_scratch/__init__.py

@@ -1,2 +1,6 @@
-def hello() -> str:
-    return "Hello from llm-agents-from-scratch!"
+# Disable the F403 warning for wildcard imports
+# ruff: noqa: F403, F401
+from .core import *
+from .core import __all__ as _core_all
+
+__all__ = sorted(_core_all)

src/llm_agents_from_scratch/base/llm.py

@@ -4,6 +4,8 @@
 
 from llm_agents_from_scratch.data_structures import ChatMessage, CompleteResult
 
+from .tool import BaseTool
+
 
 class BaseLLM(ABC):
     """Base LLM Class."""
@@ -20,11 +22,14 @@ async def complete(self, prompt: str) -> CompleteResult:
         """
 
     @abstractmethod
-    async def chat(self, chat_messages: list[ChatMessage]) -> ChatMessage:
+    async def chat(
+        self, chat_messages: list[ChatMessage], tools: list[BaseTool] = []
+    ) -> ChatMessage:
         """Chat interface.
 
         Args:
             chat_messages (list[ChatMessage]): chat history.
+            tools (list[BaseTool]): tools that the LLM can call.
 
         Returns:
             ChatMessage: The response of the LLM structured as a `ChatMessage`.

src/llm_agents_from_scratch/data_structures/llm.py

@@ -1,23 +1,42 @@
 """Data Structures for LLMs"""
 
 from enum import Enum
+from typing import Any
 
 from pydantic import BaseModel, ConfigDict
 
 
 class ChatRole(str, Enum):
+    """Roles for chat messages."""
+
     USER = "user"
     ASSISTANT = "assistant"
     SYSTEM = "system"
-    TOOL_CALL = "tool_call"
+    TOOL = "tool"
 
 
 class ChatMessage(BaseModel):
+    """The chat message data model.
+
+    Attributes:
+        role: The role of the message.
+        content: The content of the message.
+        tool_calls: Tool calls associated with the message.
+    """
+
     model_config = ConfigDict(arbitrary_types_allowed=True)
     role: ChatRole
     content: str
+    tool_calls: list[dict[str, Any]] | None = None
 
 
 class CompleteResult(BaseModel):
+    """The llm completion result data model.
+
+    Attributes:
+        response: The completion response provided by the LLM.
+        full_response: Input prompt and completion text.
+    """
+
     response: str
     full_response: str

tests/api/test_data_structure_imports.py

@@ -0,0 +1,17 @@
+import importlib
+
+import pytest
+
+from llm_agents_from_scratch.data_structures import (
+    __all__ as _data_structures_all,
+)
+
+
+@pytest.mark.parametrize("name", _data_structures_all)
+def test_data_structures_all_importable(name: str) -> None:
+    """Tests that all names listed in generators __all__ are importable."""
+    mod = importlib.import_module("llm_agents_from_scratch.data_structures")
+    attr = getattr(mod, name)
+
+    assert hasattr(mod, name)
+    assert attr is not None

tests/api/test_root_imports.py

@@ -0,0 +1,15 @@
+import importlib
+
+import pytest
+
+from llm_agents_from_scratch import __all__ as _root_all
+
+
+@pytest.mark.parametrize("name", _root_all)
+def test_root_names_all_importable(name: str) -> None:
+    """Tests that all names listed in root __all__ are importable."""
+    mod = importlib.import_module("llm_agents_from_scratch")
+    attr = getattr(mod, name)
+
+    assert hasattr(mod, name)
+    assert attr is not None

tests/conftest.py

@@ -1,6 +1,7 @@
 import pytest
 
 from llm_agents_from_scratch.base.llm import BaseLLM
+from llm_agents_from_scratch.base.tool import BaseTool
 from llm_agents_from_scratch.data_structures import ChatMessage, CompleteResult
 
 
@@ -11,7 +12,9 @@ async def complete(self, prompt: str) -> CompleteResult:
             response=result, full_response=f"{prompt} {result}"
         )
 
-    async def chat(self, chat_messages: list[ChatMessage]) -> ChatMessage:
+    async def chat(
+        self, chat_messages: list[ChatMessage], tools: list[BaseTool] = []
+    ) -> ChatMessage:
         return ChatMessage(role="assistant", content="mock chat response")
 
 

tests/llm/test_base.py

@@ -2,6 +2,7 @@
 
 
 def test_base_abstract_attr() -> None:
+    """Tests abstract methods in base class."""
     abstract_methods = BaseLLM.__abstractmethods__
 
     assert "complete" in abstract_methods