bugfixes on backport

core/just_agents/base_agent.py

@@ -1,6 +1,6 @@
 from pydantic import Field, PrivateAttr
 from typing import Optional, List, Union, Any, Generator
-from just_agents.types import Role, SupportedMessages
+from just_agents.types import MessageDict, Role, SupportedMessages
 
 from just_agents.llm_options import LLMOptions
 from just_agents.interfaces.function_call import IFunctionCall
@@ -97,7 +97,7 @@ def deepcopy_memory(self) -> IBaseMemory:
     def add_to_memory(self, messages: SupportedMessages) -> None:
         self.memory.add_message(messages)
 
-    def get_last_message(self) -> SupportedMessage:  # type: ignore
+    def get_last_message(self) -> Optional[MessageDict]:
         msg = self.memory.last_message
         if msg is None:
             raise ValueError("No messages in memory")

core/just_agents/just_tool.py

@@ -2,7 +2,7 @@
 
 from litellm.utils import function_to_dict
 from pydantic import BaseModel, Field, PrivateAttr
-from just_bus import JustEventBus
+from just_agents.just_bus import JustEventBus
 from importlib import import_module
 import inspect
 
@@ -81,8 +81,15 @@ def from_callable(cls, input_function: Callable) -> 'JustTool':
         package = input_function.__module__
         litellm_description = function_to_dict(input_function)
         arguments = cls._extract_parameters(input_function)
-
-        wrapped_callable = cls._wrap_function(input_function, litellm_description['function'])
+
+        # Get function name either from litellm description or directly from the function
+        function_name = litellm_description.get('function') or input_function.__name__
+
+        wrapped_callable = cls._wrap_function(input_function, function_name)
+
+        # Ensure function name is in litellm_description
+        if 'function' not in litellm_description:
+            litellm_description['function'] = function_name
 
         return cls(
             **litellm_description,

core/just_agents/simple/llm_session.py

@@ -9,7 +9,8 @@
 from just_agents.interfaces.agent import IAgent
 from just_agents.simple.memory import Memory
 from typing import Callable, Optional
-from just_agents.simple.streaming import AbstractStreaming
+
+from just_agents.simple.streaming.protocols.abstract_streaming import AbstractStreaming
 from just_agents.simple.streaming.openai_streaming import AsyncSession
 from just_agents.simple.utils import resolve_and_validate_agent_schema, resolve_llm_options, resolve_system_prompt, resolve_tools
 from just_agents.rotate_keys import RotateKeys

tests/test_chain_of_thought.py

@@ -2,7 +2,7 @@
 from just_agents.patterns.chain_of_throught import ChainOfThoughtAgent
 import just_agents.llm_options
 
-def count_letters(character:str, word:str) -> str:
+def count_letters(character: str, word: str) -> str:
     """ Returns the number of character occurrences in the word. """
     count:int = 0
     for char in word: