Small refactoring and adding up TBModelChat

Author: Getty

README.md

@@ -1,3 +1,18 @@
-[![Tackleberry Factory](https://raw.githubusercontent.com/Getty/tackleberry/main/tackleberry.jpg)](https://github.com/Getty/tackleberry)
+[![Tackleberry](https://raw.githubusercontent.com/Getty/tackleberry/main/tackleberry.jpg)](https://github.com/Getty/tackleberry)
 
 # Tackleberry
+
+
+
+# Install
+
+## Using PIP
+
+### Stable Version with PIP
+
+Install from `PyPi`
+
+```console
+❯ pip install --upgrade tackleberry
+```
+

tackleberry/__init__.py

@@ -1,5 +1,5 @@
 from importlib.metadata import version
-from .__main__ import TB
+from .main import TB
 
 __all__ = ['TB']
 

tackleberry/context.py

@@ -9,17 +9,22 @@ def __init__(self, content: str, role: str):
         self.content = content
 
 class TBMessageSystem(TBMessage):
-
+    """A system message"""
     def __init__(self, system_prompt: str):
         super().__init__(system_prompt, role="system")
 
-class TBMessageAssistant(TBMessage):
+class TBSystemPromptError(Exception):
+    """Error to throw if system prompt is not possible in the way as used."""
+    def __init__(self, message):
+        super().__init__(message)
 
+class TBMessageAssistant(TBMessage):
+    """An AI assistant message"""
     def __init__(self, assistant_context: str):
         super().__init__(assistant_context, role="assistant")
 
 class TBMessageUser(TBMessage):
-
+    """User message"""
     def __init__(self, user_message: str, role: Optional[str] = None):
         super().__init__(user_message, role if role is not None else "user")
 

tackleberry/engine/__init__.py

@@ -1,10 +1,20 @@
-from typing import Any, Dict, Optional
+import os
+import importlib
 
-class TBEngine:
+# Special case: Explicitly import TBEngine from base.py
+from .base import TBEngine
 
-    def __init__(self):
-        pass
+# Automatically detect all Python files in the current directory
+current_dir = os.path.dirname(__file__)
+module_files = [
+    f for f in os.listdir(current_dir)
+    if f.endswith(".py") and f not in ("__init__.py", "base.py")
+]
 
-    def model(self, model: str):
-        from ..model import TBModel
-        return TBModel(self, model)
+# Create __all__ with the module names and the special case
+__all__ = ["TBEngine"] + [os.path.splitext(f)[0] for f in module_files]
+
+# Dynamically import the modules and add them to the global namespace
+for module_name in __all__[1:]:  # Skip "TBEngine" as it's already imported
+    module = importlib.import_module(f".{module_name}", package=__name__)
+    globals()[module_name] = module

tackleberry/engine/anthropic.py

@@ -1,7 +1,6 @@
-from typing import Any, Union, Dict, List, Optional
 import os
 
-from . import TBEngine
+from .base import TBEngine
 
 class TBEngineAnthropic(TBEngine):
     default_max_tokens = 256

tackleberry/engine/base.py

@@ -0,0 +1,12 @@
+class TBEngine:
+
+    def __init__(self):
+        pass
+
+    def model(self, model: str):
+        from ..model import TBModel
+        return TBModel(self, model)
+
+    def chat(self, model: str):
+        from ..model import TBModelChat
+        return TBModelChat(self, model)

tackleberry/engine/groq.py

@@ -1,7 +1,6 @@
-from typing import Any, Union, Dict, List, Optional
 import os
 
-from . import TBEngine
+from .base import TBEngine
 
 class TBEngineGroq(TBEngine):
 

tackleberry/engine/hf.py

@@ -1,9 +1,8 @@
-from typing import Any, Union, Dict, List, Optional
 import os
 from urllib.parse import urlparse
 import base64
 
-from . import TBEngine
+from .base import TBEngine
 
 class TBEngineOllama(TBEngine):
 

tackleberry/engine/ollama.py

@@ -1,7 +1,6 @@
-from typing import Any, Union, Dict, List, Optional
 import os
 
-from . import TBEngine
+from .base import TBEngine
 
 class TBEngineOpenai(TBEngine):
 

tackleberry/engine/openai.py

@@ -0,0 +1,12 @@
+from .base import TBEngine
+
+class TBEngineTrf(TBEngine):
+
+    def __init__(self,
+        hf_token: str = None,
+        **kwargs,
+    ):
+        self.hf_token = hf_token
+
+    def __str__(self):
+        return f"TB Engine HuggingFace transformers {hex(id(self))}"

tackleberry/engine/trf.py

@@ -1,9 +1,20 @@
-from typing import Any, Dict, Optional
+import os
+import importlib
 
-from ..engine import TBEngine
+# Special case: Explicitly import TBModel from base.py
+from .base import TBModel
 
-class TBModel:
+# Automatically detect all Python files in the current directory
+current_dir = os.path.dirname(__file__)
+module_files = [
+    f for f in os.listdir(current_dir)
+    if f.endswith(".py") and f not in ("__init__.py", "base.py")
+]
 
-    def __init__(self, engine: TBEngine, model: str):
-        self.engine = engine
-        self.model = model
+# Create __all__ with the module names and the special case
+__all__ = ["TBModel"] + [os.path.splitext(f)[0] for f in module_files]
+
+# Dynamically import the modules and add them to the global namespace
+for module_name in __all__[1:]:  # Skip "TBModel" as it's already imported
+    module = importlib.import_module(f".{module_name}", package=__name__)
+    globals()[module_name] = module

tackleberry/__main__.py tackleberry/main.py

@@ -0,0 +1,9 @@
+from typing import Any, Dict, Optional
+
+from ..engine import TBEngine
+
+class TBModel:
+
+    def __init__(self, engine: TBEngine, name: str):
+        self.engine = engine
+        self.name = name

tackleberry/model/__init__.py

@@ -0,0 +1,25 @@
+from typing import Union
+
+from .base import TBModel
+from ..engine import TBEngine
+from ..context import TBContext, TBSystemPromptError
+
+class TBModelChat(TBModel):
+
+    def __init__(self,
+        engine: TBEngine,
+        model_name_or_model: Union[str|TBModel],
+        system_prompt: str = None,
+        context: TBContext = None,
+        **kwargs,
+    ):
+        self.engine = engine
+        if context is not None and system_prompt is not None:
+            raise TBSystemPromptError("A TBModelChat can't handle system_prompt and context at once.")
+        self.system_prompt = system_prompt
+        self.context = context
+        if isinstance(model_name_or_model, TBModel):
+            self.model = model_name_or_model
+        else:
+            self.model = self.engine.model(model_name_or_model)
+        self.name = model_name_or_model.name