Update tests and readme.

Author: inf3rnus

.github/workflows/workflow.yml

@@ -0,0 +1 @@
+# TODO

.gitignore

@@ -1,3 +1,7 @@
 .vscode
 **/__pycache__
-examples
+examples
+pypi_token.token
+build
+dist
+langchain_bytez.egg-info

README.md

@@ -116,6 +116,8 @@ results = bytez_chat_model_phi.invoke(messages)
 
 ```
 
+To roll our own implementation that better suites your needs, check out the implementation [here](https://github.com/Bytez-com/langchain_bytez/blob/main/langchain_bytez/BytezStdOutCallbackHandler.py)
+
 # Shutdown your cluster
 
 ```py

langchain_bytez/test.py

@@ -5,40 +5,17 @@
 
 from langchain.schema import HumanMessage, SystemMessage
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-from langchain.callbacks.stdout import StdOutCallbackHandler
 from langchain.utils import print_text
 from langchain_core.messages import AIMessage, BaseMessageChunk
-from langchain_core.outputs import LLMResult
 
 from langchain_bytez import BytezLLM
 from langchain_bytez import BytezChatModel
+from langchain_bytez import BytezStdOutCallbackHandler
 
 
 # sudo cat /var/log/cloud-init-output.log
 
 
-class ExtendedStdOutCallbackHandler(StdOutCallbackHandler):
-    def print(self, *args, color: str = None, **kwargs):
-        text = " ".join([str(arg) for arg in args])
-        print_text(f"\n{text}", color or self.color, **kwargs)
-
-    def _on_start(self, serialized, prompts, *args, **kwargs):
-        model_type = kwargs["invocation_params"]["model_type"]
-        model_name = kwargs["invocation_params"]["model_name"]
-
-        self.print(f"Model started: {model_name} ({model_type})", color="blue")
-        self.print("Prompt:", prompts, color="yellow")
-
-    def on_llm_start(self, *args, **kwargs):
-        self._on_start(*args, **kwargs)
-
-    def on_llm_end(self, response: LLMResult, *args, **kwargs):
-        self.print("Model responded with: ", response, "\n")
-
-    def on_chat_model_start(self, *args, **kwargs):
-        self._on_start(*args, **kwargs)
-
-
 async def run_async_test():
     API_KEY = os.environ.get("API_KEY")
 
@@ -61,7 +38,7 @@ async def run_async_test():
     callbacks = [
         #
         StreamingStdOutCallbackHandler(),
-        ExtendedStdOutCallbackHandler(color="green"),
+        BytezStdOutCallbackHandler(color="green"),
     ]
 
     bytez_chat_model_phi = BytezChatModel(

setup.py

@@ -2,13 +2,13 @@
 
 setup(
     name="langchain_bytez",  # Name of your package
-    version="0.0.0",
+    version="0.0.5",
     packages=find_packages(),
     description="Bytez langchain integration",
     long_description=open("./README.md").read(),
     long_description_content_type="text/markdown",
     author="Bytez",
-    url="https://github.com/bytez-com/docs",
+    url="https://github.com/Bytez-com/langchain_bytez",
     install_requires=["langchain==0.3.17"],
     classifiers=[
         "Development Status :: 4 - Beta",