Merge pull request #399 from MervinPraison/develop

Bump version to 2.0.80 and update praisonaiagents dependency
2026-11-07 17:33:33 -10:00 · 2025-02-28 10:25:47 +00:00 · 2025-02-28 10:25:32 +00:00 · 2025-02-28 07:26:52 +00:00 · 2025-02-28 07:26:22 +00:00 · 2025-02-28 07:05:13 +00:00
25 changed files with 1691 additions and 240 deletions
@@ -118,6 +118,8 @@ const agent = new Agent({ instructions: 'You are a helpful AI assistant' });
 agent.start('Write a movie script about a robot in Mars');
 ```

+![PraisonAI CLI Demo](docs/demo/praisonai-cli-demo.gif)
+
 ## AI Agents Flow

 ```mermaid
@@ -1,6 +1,6 @@
 FROM python:3.11-slim
 WORKDIR /app
 COPY . .
-RUN pip install flask praisonai==2.0.62 gunicorn markdown
+RUN pip install flask praisonai==2.0.80 gunicorn markdown
 EXPOSE 8080
 CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]
@@ -110,7 +110,7 @@ Loads environment variables from .env file or system and sets them.</p>
            file.write(&#34;FROM python:3.11-slim\n&#34;)
            file.write(&#34;WORKDIR /app\n&#34;)
            file.write(&#34;COPY . .\n&#34;)
-            file.write(&#34;RUN pip install flask praisonai==2.0.62 gunicorn markdown\n&#34;)
+            file.write(&#34;RUN pip install flask praisonai==2.0.80 gunicorn markdown\n&#34;)
            file.write(&#34;EXPOSE 8080\n&#34;)
            file.write(&#39;CMD [&#34;gunicorn&#34;, &#34;-b&#34;, &#34;0.0.0.0:8080&#34;, &#34;api:app&#34;]\n&#39;)
            
@@ -0,0 +1,117 @@
+---
+title: "Command Line Interface"
+sidebarTitle: "CLI"
+description: "Use PraisonAI directly from your terminal with simple commands"
+icon: "terminal"
+---
+
+```mermaid
+flowchart LR
+    Start((🚀)) --> CLI[CLI Command]
+    CLI --> Agent[AI Agent]
+    Agent --> Config[Config/YAML]
+    Config --> Out((🏁))
+    
+    style Start fill:#8B0000,color:#fff,stroke:#8B0000
+    style CLI fill:#2E8B57,color:#fff
+    style Agent fill:#2E8B57,color:#fff
+    style Config fill:#2E8B57,color:#fff
+    style Out fill:#8B0000,color:#fff,stroke:#8B0000
+```
+
+PraisonAI CLI provides a simple way to interact with AI agents directly from your terminal. You can run quick commands, specify LLM options, or use YAML configuration files for more complex scenarios.
+
+<Frame>
+  <img src="../demo/praisonai-cli-demo.gif" alt="PraisonAI CLI Demo" />
+</Frame>
+
+## Quick Start
+
+<Steps>
+  <Step title="Install Package">
+    Install the PraisonAI package:
+    ```bash
+    pip install praisonai
+    ```
+  </Step>
+
+  <Step title="Set API Key">
+    Set your OpenAI API key as an environment variable:
+    ```bash
+    export OPENAI_API_KEY=your_api_key_here
+    ```
+  </Step>
+</Steps>
+
+## Usage Examples
+
+<AccordionGroup>
+  <Accordion title="Simple Command" icon="terminal" defaultOpen>
+    Run a simple command directly:
+    ```bash
+    praisonai "write a movie script in 3 lines"
+    ```
+  </Accordion>
+
+  <Accordion title="With LLM Option" icon="wand-magic-sparkles" defaultOpen>
+    Specify a different LLM model:
+    ```bash
+    praisonai "write a movie script in 3 lines" --llm gpt-4o-mini
+    ```
+  </Accordion>
+
+  <Accordion title="Using YAML Config" icon="file-code" defaultOpen>
+    Run agents defined in a YAML file:
+    ```bash
+    praisonai agents.yaml
+    ```
+  </Accordion>
+</AccordionGroup>
+
+## Configuration
+
+<Tabs>
+  <Tab title="Create Config">
+    Initialize a new agents.yaml file for your project:
+    ```bash
+    praisonai --init "Create a movie script about AI"
+    ```
+    This will create an `agents.yaml` file with predefined configuration for your task.
+  </Tab>
+
+  <Tab title="Auto Mode">
+    Let PraisonAI automatically create and configure agents:
+    ```bash
+    praisonai --auto "Create a movie script about AI"
+    ```
+    This mode will analyze your task and set up appropriate agents automatically.
+  </Tab>
+</Tabs>
+
+## Features
+
+<CardGroup cols={2}>
+  <Card title="Simple Commands" icon="terminal">
+    Run AI tasks directly from your terminal with simple commands.
+  </Card>
+  <Card title="LLM Options" icon="wand-magic-sparkles">
+    Choose from different LLM models for your specific needs.
+  </Card>
+  <Card title="YAML Support" icon="file-code">
+    Use YAML files for complex agent configurations and workflows.
+  </Card>
+  <Card title="Auto Configuration" icon="robot">
+    Automatic agent setup based on task requirements.
+  </Card>
+</CardGroup>
+
+## Next Steps
+
+<CardGroup>
+  <Card title="YAML Configuration" icon="file-code" href="/config">
+    Learn more about YAML configuration options
+  </Card>
+  <Card title="API Reference" icon="code" href="/api-reference">
+    View the complete API documentation
+  </Card>
+</CardGroup>
@@ -119,6 +119,7 @@
    {
      "group": "Features",
      "pages": [
+        "features/cli",
        "features/autoagents",
        "features/image-generation",
        "features/selfreflection",
@@ -20,6 +20,20 @@ icon: "dumbbell"
 export HF_TOKEN=xxxxxxxxxxxxxx
 ```

+## Initilise praisonai train
+
+```bash
+praisonai train init
+```
+
+## Required
+
+1. Huggingface token
+2. Base model to train on (e.g. unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit)
+3. Dataset to train on (e.g. yahma/alpaca-cleaned)
+4. Huggingface model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)
+5. Ollama model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)
+
 ## To upload to ollama.com (Linux)

 ```bash
@@ -108,4 +122,21 @@ quantization_method:

 ```bash
 praisonai train
+```
+
+## wandb
+
+```bash
+wandb login
+```
+
+<Note>
+Get the key from [here](https://wandb.ai/site/login)
+</Note>
+
+```bash
+export PRAISON_WANDB=True
+export WANDB_LOG_MODEL=checkpoint
+export WANDB_PROJECT=praisonai-test
+export PRAISON_WANDB_RUN_NAME=praisonai-train  
 ```
@@ -3,7 +3,7 @@ class Praisonai < Formula
  
    desc "AI tools for various AI applications"
    homepage "https://github.com/MervinPraison/PraisonAI"
-    url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/2.0.62.tar.gz"
+    url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/2.0.80.tar.gz"
    sha256 "1828fb9227d10f991522c3f24f061943a254b667196b40b1a3e4a54a8d30ce32"  # Replace with actual SHA256 checksum
    license "MIT"
  
@@ -25,7 +25,6 @@ CALL_MODULE_AVAILABLE = False
 CREWAI_AVAILABLE = False
 AUTOGEN_AVAILABLE = False
 PRAISONAI_AVAILABLE = False
-
 try:
    # Create necessary directories and set CHAINLIT_APP_ROOT
    if "CHAINLIT_APP_ROOT" not in os.environ:
@@ -140,6 +139,8 @@ class PraisonAI:
        provided arguments.
        """
        args = self.parse_args()
+        # Store args for use in handle_direct_prompt
+        self.args = args
        invocation_cmd = "praisonai"
        version_string = f"PraisonAI version {__version__}"

@@ -188,8 +189,7 @@ class PraisonAI:
            package_root = os.path.dirname(os.path.abspath(__file__))
            config_yaml_destination = os.path.join(os.getcwd(), 'config.yaml')

-            # Create config.yaml only if it doesn't exist or --model or --dataset is provided
-            if not os.path.exists(config_yaml_destination) or args.model or args.dataset:
+            if not os.path.exists(config_yaml_destination):
                config = generate_config(
                    model_name=args.model,
                    hf_model_name=args.hf,
@@ -200,6 +200,20 @@ class PraisonAI:
                )
                with open('config.yaml', 'w') as f:
                    yaml.dump(config, f, default_flow_style=False, indent=2)
+            elif args.model or args.hf or args.ollama or (args.dataset and args.dataset != "yahma/alpaca-cleaned"):
+                config = generate_config(
+                    model_name=args.model,
+                    hf_model_name=args.hf,
+                    ollama_model_name=args.ollama,
+                    dataset=[{
+                        "name": args.dataset
+                    }]
+                )
+                with open('config.yaml', 'w') as f:
+                    yaml.dump(config, f, default_flow_style=False, indent=2)
+            else:
+                with open(config_yaml_destination, 'r') as f:
+                    config = yaml.safe_load(f)

            # Overwrite huggingface_save and ollama_save if --hf or --ollama are provided
            if args.hf:
@@ -226,7 +240,18 @@ class PraisonAI:
                print("All packages installed.")

            train_args = sys.argv[2:]  # Get all arguments after 'train'
-            train_script_path = os.path.join(package_root, 'train.py')
+            
+            # Check if this is a vision model - handle all case variations
+            model_name = config.get("model_name", "").lower()
+            is_vision_model = any(x in model_name for x in ["-vl-", "-vl", "vl-", "-vision-", "-vision", "vision-", "visionmodel"])
+            
+            # Choose appropriate training script
+            if is_vision_model:
+                train_script_path = os.path.join(package_root, 'train_vision.py')
+                print("Using vision training script for VL model...")
+            else:
+                train_script_path = os.path.join(package_root, 'train.py')
+                print("Using standard training script...")

            # Set environment variables
            env = os.environ.copy()
@@ -304,6 +329,7 @@ class PraisonAI:
        parser.add_argument("command", nargs="?", help="Command to run or direct prompt")
        parser.add_argument("--deploy", action="store_true", help="Deploy the application")
        parser.add_argument("--model", type=str, help="Model name")
+        parser.add_argument("--llm", type=str, help="LLM model to use for direct prompts")
        parser.add_argument("--hf", type=str, help="Hugging Face model name")
        parser.add_argument("--ollama", type=str, help="Ollama model name")
        parser.add_argument("--dataset", type=str, help="Dataset name for training", default="yahma/alpaca-cleaned")
@@ -390,9 +416,8 @@ class PraisonAI:
                sys.exit(0)

            elif args.command == 'train':
-                print("[red]ERROR: Train feature is not installed. Install with:[/red]")
-                print("\npip install \"praisonai[train]\"\n")
-                sys.exit(1)
+                package_root = os.path.dirname(os.path.abspath(__file__))
+                config_yaml_destination = os.path.join(os.getcwd(), 'config.yaml')

            elif args.command == 'ui':
                if not CHAINLIT_AVAILABLE:
@@ -424,21 +449,33 @@ class PraisonAI:
        Handle direct prompt by creating a single agent and running it.
        """
        if PRAISONAI_AVAILABLE:
-            agent = PraisonAgent(
-                name="DirectAgent",
-                role="Assistant",
-                goal="Complete the given task",
-                backstory="You are a helpful AI assistant"
-            )
-            agent.start(prompt)
+            agent_config = {
+                "name": "DirectAgent",
+                "role": "Assistant",
+                "goal": "Complete the given task",
+                "backstory": "You are a helpful AI assistant"
+            }
+            
+            # Add llm if specified
+            if hasattr(self, 'args') and self.args.llm:
+                agent_config["llm"] = self.args.llm
+            
+            agent = PraisonAgent(**agent_config)
+            result = agent.start(prompt)
            return ""
        elif CREWAI_AVAILABLE:
-            agent = Agent(
-                name="DirectAgent",
-                role="Assistant",
-                goal="Complete the given task",
-                backstory="You are a helpful AI assistant"
-            )
+            agent_config = {
+                "name": "DirectAgent",
+                "role": "Assistant",
+                "goal": "Complete the given task",
+                "backstory": "You are a helpful AI assistant"
+            }
+            
+            # Add llm if specified
+            if hasattr(self, 'args') and self.args.llm:
+                agent_config["llm"] = self.args.llm
+            
+            agent = Agent(**agent_config)
            task = Task(
                description=prompt,
                agent=agent
@@ -450,6 +487,10 @@ class PraisonAI:
            return crew.kickoff()
        elif AUTOGEN_AVAILABLE:
            config_list = self.config_list
+            # Add llm if specified
+            if hasattr(self, 'args') and self.args.llm:
+                config_list[0]['model'] = self.args.llm
+                
            assistant = autogen.AssistantAgent(
                name="DirectAgent",
                llm_config={"config_list": config_list}
@@ -56,7 +56,7 @@ class CloudDeployer:
            file.write("FROM python:3.11-slim\n")
            file.write("WORKDIR /app\n")
            file.write("COPY . .\n")
-            file.write("RUN pip install flask praisonai==2.0.62 gunicorn markdown\n")
+            file.write("RUN pip install flask praisonai==2.0.80 gunicorn markdown\n")
            file.write("EXPOSE 8080\n")
            file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
            
@@ -11,6 +11,38 @@ if [[ "$OSTYPE" == "darwin"* ]]; then
 elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
    # Linux
    MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh"
+    
+    # Install libcurl development package if not present (Debian based)
+    if command -v dpkg &> /dev/null; then
+        if ! dpkg -s libcurl4-openssl-dev &> /dev/null; then
+            echo "libcurl4-openssl-dev is not installed. Installing..."
+            sudo apt-get update
+            sudo apt-get install -y libcurl4-openssl-dev
+        else
+            echo "libcurl4-openssl-dev is already installed."
+        fi
+    else
+        echo "Non-Debian based Linux detected. Please ensure libcurl development libraries are installed."
+    fi
+
+    # Check if ollama is installed and executable; if not, install it
+    if ! command -v ollama &> /dev/null; then
+        echo "Ollama is not installed. Installing Ollama..."
+        curl -fsSL https://ollama.com/install.sh | sh
+        
+        # Generate SSH key non-interactively only if it doesn't already exist
+        if [ ! -f ~/.ssh/id_ed25519 ]; then
+            echo "Generating SSH key for Ollama..."
+            ssh-keygen -t ed25519 -N "" -f ~/.ssh/id_ed25519 -q
+        else
+            echo "SSH key ~/.ssh/id_ed25519 already exists. Skipping generation."
+        fi
+        echo "Copying SSH key to /usr/share/ollama/.ollama..."
+        sudo cp ~/.ssh/id_ed25519 /usr/share/ollama/.ollama
+    else
+        echo "Ollama is already installed."
+    fi
+
 elif [[ "$OSTYPE" == "msys" || "$OSTYPE" == "win32" ]]; then
    # Windows
    MINICONDA_URL="https://repo.anaconda.com/miniconda/Miniconda3-latest-Windows-x86_64.exe"
@@ -21,12 +53,12 @@ else
    exit 1
 fi

-# Check if conda is already installed
+# Check if conda is installed
 if ! command -v conda &> /dev/null; then
    echo "Conda is not installed. Installing Miniconda..."
-    wget $MINICONDA_URL -O ~/miniconda.sh
-    bash ~/miniconda.sh -b -p $HOME/miniconda
-    source $HOME/miniconda/bin/activate
+    wget "$MINICONDA_URL" -O ~/miniconda.sh
+    bash ~/miniconda.sh -b -p "$HOME/miniconda"
+    source "$HOME/miniconda/bin/activate"
    conda init
 else
    echo "Conda is already installed."
@@ -34,39 +66,40 @@ fi

 # Create and activate the Conda environment
 ENV_NAME="praison_env"
-if conda info --envs | grep -q $ENV_NAME; then
+if conda info --envs | grep -q "$ENV_NAME"; then
    echo "Environment $ENV_NAME already exists. Recreating..."
-    conda env remove -y -n $ENV_NAME  # Remove existing environment
+    conda env remove -y -n "$ENV_NAME"
    if [[ "$OSTYPE" == "darwin"* ]]; then
-        # macOS (both Intel and M1/M2)
-        conda create --name $ENV_NAME python=3.10 pytorch=2.3.0 -c pytorch -y
+        conda create --name "$ENV_NAME" python=3.10 pytorch=2.3.0 -c pytorch -y
    elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
-        # Linux
-        conda create --name $ENV_NAME python=3.10 pytorch=2.3.0 cudatoolkit=11.8 -c pytorch -c nvidia -y
+        conda create --name "$ENV_NAME" python=3.10 pytorch=2.3.0 cudatoolkit=11.8 -c pytorch -c nvidia -y
    fi
-    # conda activate $ENV_NAME
 else
    echo "Creating new environment $ENV_NAME..."
    if [[ "$OSTYPE" == "darwin"* ]]; then
-        # macOS (both Intel and M1/M2)
-        conda create --name $ENV_NAME python=3.10 pytorch=2.3.0 -c pytorch -y
+        conda create --name "$ENV_NAME" python=3.10 pytorch=2.3.0 -c pytorch -y
    elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
-        # Linux
-        conda create --name $ENV_NAME python=3.10 pytorch=2.3.0 cudatoolkit=11.8 -c pytorch -c nvidia -y
+        conda create --name "$ENV_NAME" python=3.10 pytorch=2.3.0 cudatoolkit=11.8 -c pytorch -c nvidia -y
    fi
-    # conda activate $ENV_NAME
 fi

-# source $HOME/miniconda/bin/activate $ENV_NAME
+# Activate the environment
+source "$HOME/miniconda/bin/activate" "$ENV_NAME"

-# Get full path of pip
-PIP_FULL_PATH=$(conda run -n $ENV_NAME which pip)
+# Install cmake via conda
+echo "Installing cmake..."
+conda install -y cmake

-# Install other packages within the activated environment
-# Use PIP_FULL_PATH to run pip commands
+# Get full path of pip within the activated environment
+PIP_FULL_PATH=$(conda run -n "$ENV_NAME" which pip)
+
+# Install other packages using pip
 $PIP_FULL_PATH install --upgrade pip 
 $PIP_FULL_PATH install "xformers==0.0.26.post1"
-$PIP_FULL_PATH install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git@4e570be9ae4ced8cdc64e498125708e34942befc"
+$PIP_FULL_PATH install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git@53a773e4fbc53a1d96c7ba107e5fe75dab07027b"
 $PIP_FULL_PATH install --no-deps "trl<0.9.0" peft accelerate bitsandbytes 
+$PIP_FULL_PATH install unsloth_zoo
+$PIP_FULL_PATH install cut_cross_entropy
+$PIP_FULL_PATH install sentencepiece protobuf datasets huggingface_hub hf_transfer wandb

 echo "Setup completed successfully!"
@@ -1,168 +1,287 @@
-import subprocess
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+This script finetunes a model using Unsloth's fast training framework.
+It supports both ShareGPT and Alpaca‑style datasets by converting raw conversation
+data into plain-text prompts using a chat template, then pre‑tokenizing the prompts.
+Extra debug logging is added to help trace the root cause of errors.
+"""
+
 import os
 import sys
 import yaml
 import torch
 import shutil
+import subprocess
 from transformers import TextStreamer
 from unsloth import FastLanguageModel, is_bfloat16_supported
 from trl import SFTTrainer
 from transformers import TrainingArguments
-from datasets import load_dataset, concatenate_datasets, Dataset
+from datasets import load_dataset, concatenate_datasets
 from psutil import virtual_memory
+from unsloth.chat_templates import standardize_sharegpt, get_chat_template
+from functools import partial

-class train:
+#####################################
+# Step 1: Formatting Raw Conversations
+#####################################
+def formatting_prompts_func(examples, tokenizer):
+    """
+    Converts each example's conversation into a single plain-text prompt.
+    If the example has a "conversations" field, process it as ShareGPT-style.
+    Otherwise, assume Alpaca-style data with "instruction", "input", and "output" fields.
+    """
+    print("DEBUG: formatting_prompts_func() received batch with keys:", list(examples.keys()))
+    texts = []
+    # Check if the example has a "conversations" field.
+    if "conversations" in examples:
+        for convo in examples["conversations"]:
+            try:
+                formatted = tokenizer.apply_chat_template(
+                    convo,
+                    tokenize=False,  # Return a plain string
+                    add_generation_prompt=False
+                )
+            except Exception as e:
+                print(f"ERROR in apply_chat_template (conversations): {e}")
+                formatted = ""
+            # Flatten list if necessary
+            if isinstance(formatted, list):
+                formatted = formatted[0] if len(formatted) == 1 else "\n".join(formatted)
+            texts.append(formatted)
+    else:
+        # Assume Alpaca format: use "instruction", "input", and "output" keys.
+        instructions = examples.get("instruction", [])
+        inputs_list = examples.get("input", [])
+        outputs_list = examples.get("output", [])
+        # If any field is missing, replace with empty string.
+        for ins, inp, out in zip(instructions, inputs_list, outputs_list):
+            # Create a conversation-like structure.
+            convo = [
+                {"role": "user", "content": ins + (f"\nInput: {inp}" if inp.strip() != "" else "")},
+                {"role": "assistant", "content": out}
+            ]
+            try:
+                formatted = tokenizer.apply_chat_template(
+                    convo,
+                    tokenize=False,
+                    add_generation_prompt=False
+                )
+            except Exception as e:
+                print(f"ERROR in apply_chat_template (alpaca): {e}")
+                formatted = ""
+            if isinstance(formatted, list):
+                formatted = formatted[0] if len(formatted) == 1 else "\n".join(formatted)
+            texts.append(formatted)
+    if texts:
+        print("DEBUG: Raw texts sample (first 200 chars):", texts[0][:200])
+    return {"text": texts}
+
+#####################################
+# Step 2: Tokenizing the Prompts
+#####################################
+def tokenize_function(examples, hf_tokenizer, max_length):
+    """
+    Tokenizes a batch of text prompts with padding and truncation enabled.
+    """
+    flat_texts = []
+    for t in examples["text"]:
+        if isinstance(t, list):
+            t = t[0] if len(t) == 1 else " ".join(t)
+        flat_texts.append(t)
+    print("DEBUG: Tokenizing a batch of size:", len(flat_texts))
+    tokenized = hf_tokenizer(
+        flat_texts,
+        padding="max_length",
+        truncation=True,
+        max_length=max_length,
+        return_tensors="pt",
+    )
+    tokenized = {key: value.tolist() for key, value in tokenized.items()}
+    sample_key = list(tokenized.keys())[0]
+    print("DEBUG: Tokenized sample (first 10 tokens of", sample_key, "):", tokenized[sample_key][0][:10])
+    return tokenized
+
+#####################################
+# Main Training Class
+#####################################
+class TrainModel:
    def __init__(self, config_path="config.yaml"):
        self.load_config(config_path)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model, self.tokenizer = None, None
+        self.model = None
+        self.hf_tokenizer = None   # The underlying HF tokenizer
+        self.chat_tokenizer = None # Chat wrapper for formatting

    def load_config(self, path):
        with open(path, "r") as file:
            self.config = yaml.safe_load(file)
+        print("DEBUG: Loaded config:", self.config)

    def print_system_info(self):
-        print(f"PyTorch version: {torch.__version__}")
-        print(f"CUDA version: {torch.version.cuda}")
+        print("DEBUG: PyTorch version:", torch.__version__)
+        print("DEBUG: CUDA version:", torch.version.cuda)
        if torch.cuda.is_available():
-            device_capability = torch.cuda.get_device_capability()
-            print(f"CUDA Device Capability: {device_capability}")
+            print("DEBUG: CUDA Device Capability:", torch.cuda.get_device_capability())
        else:
-            print("CUDA is not available")
-
-        python_version = sys.version
-        pip_version = subprocess.check_output(['pip', '--version']).decode().strip()
-        python_path = sys.executable
-        pip_path = subprocess.check_output(['which', 'pip']).decode().strip()
-        print(f"Python Version: {python_version}")
-        print(f"Pip Version: {pip_version}")
-        print(f"Python Path: {python_path}")
-        print(f"Pip Path: {pip_path}")
+            print("DEBUG: CUDA is not available")
+        print("DEBUG: Python Version:", sys.version)
+        print("DEBUG: Python Path:", sys.executable)

    def check_gpu(self):
        gpu_stats = torch.cuda.get_device_properties(0)
-        print(f"GPU = {gpu_stats.name}. Max memory = {round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)} GB.")
+        print(f"DEBUG: GPU = {gpu_stats.name}. Max memory = {round(gpu_stats.total_memory/(1024**3),3)} GB.")

    def check_ram(self):
        ram_gb = virtual_memory().total / 1e9
-        print('Your runtime has {:.1f} gigabytes of available RAM\n'.format(ram_gb))
+        print(f"DEBUG: Your runtime has {ram_gb:.1f} gigabytes of available RAM")
        if ram_gb < 20:
-            print('Not using a high-RAM runtime')
+            print("DEBUG: Not using a high-RAM runtime")
        else:
-            print('You are using a high-RAM runtime!')
-
-    # def install_packages(self):
-    #     subprocess.run(["pip", "install", "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git@4e570be9ae4ced8cdc64e498125708e34942befc"])
-    #     subprocess.run(["pip", "install", "--no-deps", "trl<0.9.0", "peft==0.12.0", "accelerate==0.33.0", "bitsandbytes==0.43.3"])
+            print("DEBUG: You are using a high-RAM runtime!")

    def prepare_model(self):
-        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
+        print("DEBUG: Preparing model and tokenizer...")
+        self.model, original_tokenizer = FastLanguageModel.from_pretrained(
            model_name=self.config["model_name"],
            max_seq_length=self.config["max_seq_length"],
            dtype=None,
-            load_in_4bit=self.config["load_in_4bit"]
+            load_in_4bit=self.config["load_in_4bit"],
        )
+        print("DEBUG: Model and original tokenizer loaded.")
+        if original_tokenizer.pad_token is None:
+            original_tokenizer.pad_token = original_tokenizer.eos_token
+        original_tokenizer.model_max_length = self.config["max_seq_length"]
+        self.chat_tokenizer = get_chat_template(original_tokenizer, chat_template="llama-3.1")
+        self.hf_tokenizer = original_tokenizer
+        print("DEBUG: Chat tokenizer created; HF tokenizer saved.")
        self.model = FastLanguageModel.get_peft_model(
            self.model,
-            r=self.config["lora_r"],
-            target_modules=self.config["lora_target_modules"],
-            lora_alpha=self.config["lora_alpha"],
-            lora_dropout=self.config["lora_dropout"],
-            bias=self.config["lora_bias"],
-            use_gradient_checkpointing=self.config["use_gradient_checkpointing"],
-            random_state=self.config["random_state"],
-            use_rslora=self.config["use_rslora"],
-            loftq_config=self.config["loftq_config"],
+            r=16,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+            lora_alpha=16,
+            lora_dropout=0,
+            bias="none",
+            use_gradient_checkpointing="unsloth",
+            random_state=3407,
+            use_rslora=False,
+            loftq_config=None,
        )
+        print("DEBUG: LoRA adapters added.")

    def process_dataset(self, dataset_info):
        dataset_name = dataset_info["name"]
        split_type = dataset_info.get("split_type", "train")
-        processing_func = getattr(self, dataset_info.get("processing_func", "format_prompts"))
-        rename = dataset_info.get("rename", {})
-        filter_data = dataset_info.get("filter_data", False)
-        filter_column_value = dataset_info.get("filter_column_value", "id")
-        filter_value = dataset_info.get("filter_value", "alpaca")
-        num_samples = dataset_info.get("num_samples", 20000)
-
+        print(f"DEBUG: Loading dataset '{dataset_name}' split '{split_type}'...")
        dataset = load_dataset(dataset_name, split=split_type)
-
-        if rename:
-            dataset = dataset.rename_columns(rename)
-        if filter_data:
-            dataset = dataset.filter(lambda example: filter_value in example[filter_column_value]).shuffle(seed=42).select(range(num_samples))
-        dataset = dataset.map(processing_func, batched=True)
+        print("DEBUG: Dataset columns:", dataset.column_names)
+        if "conversations" in dataset.column_names:
+            print("DEBUG: Standardizing dataset (ShareGPT style)...")
+            dataset = standardize_sharegpt(dataset)
+        else:
+            print("DEBUG: Dataset does not have 'conversations'; assuming Alpaca format.")
+        print("DEBUG: Applying formatting function to dataset...")
+        format_func = partial(formatting_prompts_func, tokenizer=self.chat_tokenizer)
+        dataset = dataset.map(format_func, batched=True, remove_columns=dataset.column_names)
+        sample = dataset[0]
+        print("DEBUG: Sample processed example keys:", list(sample.keys()))
+        if "text" in sample:
+            print("DEBUG: Sample processed 'text' type:", type(sample["text"]))
+            print("DEBUG: Sample processed 'text' content (first 200 chars):", sample["text"][:200])
+        else:
+            print("DEBUG: Processed sample does not contain 'text'.")
        return dataset

-    def format_prompts(self, examples):
-        alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-
-        ### Instruction:
-        {}
-
-        ### Input:
-        {}
-
-        ### Response:
-        {}"""
-        texts = [alpaca_prompt.format(ins, inp, out) + self.tokenizer.eos_token for ins, inp, out in zip(examples["instruction"], examples["input"], examples["output"])]
-        return {"text": texts}
+    def tokenize_dataset(self, dataset):
+        print("DEBUG: Tokenizing the entire dataset...")
+        tokenized_dataset = dataset.map(
+            lambda examples: tokenize_function(examples, self.hf_tokenizer, self.config["max_seq_length"]),
+            batched=True
+        )
+        tokenized_dataset = tokenized_dataset.remove_columns(["text"])
+        print("DEBUG: Tokenized dataset sample keys:", tokenized_dataset[0].keys())
+        return tokenized_dataset

    def load_datasets(self):
        datasets = []
        for dataset_info in self.config["dataset"]:
+            print("DEBUG: Processing dataset info:", dataset_info)
            datasets.append(self.process_dataset(dataset_info))
-        return concatenate_datasets(datasets)
+        combined = concatenate_datasets(datasets)
+        print("DEBUG: Combined dataset has", len(combined), "examples.")
+        return combined

    def train_model(self):
-        dataset = self.load_datasets()
+        print("DEBUG: Starting training...")
+        raw_dataset = self.load_datasets()
+        tokenized_dataset = self.tokenize_dataset(raw_dataset)
+        print("DEBUG: Dataset tokenization complete.")
+        # Build the training arguments parameters dynamically
+        ta_params = {
+            "per_device_train_batch_size": self.config.get("per_device_train_batch_size", 2),
+            "gradient_accumulation_steps": self.config.get("gradient_accumulation_steps", 2),
+            "warmup_steps": self.config.get("warmup_steps", 50),
+            "max_steps": self.config.get("max_steps", 2800),
+            "learning_rate": self.config.get("learning_rate", 2e-4),
+            "fp16": self.config.get("fp16", not is_bfloat16_supported()),
+            "bf16": self.config.get("bf16", is_bfloat16_supported()),
+            "logging_steps": self.config.get("logging_steps", 15),
+            "optim": self.config.get("optim", "adamw_8bit"),
+            "weight_decay": self.config.get("weight_decay", 0.01),
+            "lr_scheduler_type": self.config.get("lr_scheduler_type", "linear"),
+            "seed": self.config.get("seed", 3407),
+            "output_dir": self.config.get("output_dir", "outputs"),
+            "report_to": "none" if not os.getenv("PRAISON_WANDB") else "wandb",
+            "remove_unused_columns": self.config.get("remove_unused_columns", False)
+        }
+        if os.getenv("PRAISON_WANDB"):
+            ta_params["save_steps"] = self.config.get("save_steps", 100)
+            ta_params["run_name"] = os.getenv("PRAISON_WANDB_RUN_NAME", "praisonai-train")
+
+        training_args = TrainingArguments(**ta_params)
+        # Since the dataset is pre-tokenized, we supply a dummy dataset_text_field.
        trainer = SFTTrainer(
            model=self.model,
-            tokenizer=self.tokenizer,
-            train_dataset=dataset,
-            dataset_text_field=self.config["dataset_text_field"],
+            tokenizer=self.hf_tokenizer,
+            train_dataset=tokenized_dataset,
+            dataset_text_field="input_ids",  # Dummy field since data is numeric
            max_seq_length=self.config["max_seq_length"],
-            dataset_num_proc=self.config["dataset_num_proc"],
-            packing=self.config["packing"],
-            args=TrainingArguments(
-                per_device_train_batch_size=self.config["per_device_train_batch_size"],
-                gradient_accumulation_steps=self.config["gradient_accumulation_steps"],
-                warmup_steps=self.config["warmup_steps"],
-                num_train_epochs=self.config["num_train_epochs"],
-                max_steps=self.config["max_steps"],
-                learning_rate=self.config["learning_rate"],
-                fp16=not is_bfloat16_supported(),
-                bf16=is_bfloat16_supported(),
-                logging_steps=self.config["logging_steps"],
-                optim=self.config["optim"],
-                weight_decay=self.config["weight_decay"],
-                lr_scheduler_type=self.config["lr_scheduler_type"],
-                seed=self.config["seed"],
-                output_dir=self.config["output_dir"],
-            ),
+            dataset_num_proc=1,  # Use a single process to avoid pickling issues
+            packing=False,
+            args=training_args,
        )
+        from unsloth.chat_templates import train_on_responses_only
+        trainer = train_on_responses_only(
+            trainer,
+            instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
+            response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
+        )
+        print("DEBUG: Beginning trainer.train() ...")
        trainer.train()
-        self.model.save_pretrained("lora_model") # Local saving
-        self.tokenizer.save_pretrained("lora_model")
+        print("DEBUG: Training complete. Saving model and tokenizer locally...")
+        self.model.save_pretrained("lora_model")
+        self.hf_tokenizer.save_pretrained("lora_model")
+        print("DEBUG: Saved model and tokenizer to 'lora_model'.")

    def inference(self, instruction, input_text):
        FastLanguageModel.for_inference(self.model)
-        alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+        messages = [{"role": "user", "content": f"{instruction}\n\nInput: {input_text}"}]
+        inputs = self.hf_tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt"
+        ).to("cuda")
+        outputs = self.model.generate(
+            input_ids=inputs,
+            max_new_tokens=64,
+            use_cache=True,
+            temperature=1.5,
+            min_p=0.1
+        )
+        print("DEBUG: Inference output:", self.hf_tokenizer.batch_decode(outputs))

-        ### Instruction:
-        {}
-
-        ### Input:
-        {}
-
-        ### Response:
-        {}"""
-        inputs = self.tokenizer([alpaca_prompt.format(instruction, input_text, "")], return_tensors="pt").to("cuda")
-        outputs = self.model.generate(**inputs, max_new_tokens=64, use_cache=True)
-        print(self.tokenizer.batch_decode(outputs))
-        
    def load_model(self):
-        """Loads the model and tokenizer using the FastLanguageModel library."""
        from unsloth import FastLanguageModel
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=self.config["output_dir"],
@@ -177,56 +296,235 @@ class train:
            shutil.rmtree(self.config["hf_model_name"])
        self.model.push_to_hub_merged(
            self.config["hf_model_name"],
-            self.tokenizer,
+            self.hf_tokenizer,
            save_method="merged_16bit",
-            token=os.getenv('HF_TOKEN')
+            token=os.getenv("HF_TOKEN")
        )

    def push_model_gguf(self):
        self.model.push_to_hub_gguf(
            self.config["hf_model_name"],
-            self.tokenizer,
+            self.hf_tokenizer,
            quantization_method=self.config["quantization_method"],
-            token=os.getenv('HF_TOKEN')
+            token=os.getenv("HF_TOKEN")
        )
-    
+
    def save_model_gguf(self):
        self.model.save_pretrained_gguf(
            self.config["hf_model_name"],
-            self.tokenizer,
+            self.hf_tokenizer,
            quantization_method="q4_k_m"
        )

    def prepare_modelfile_content(self):
        output_model = self.config["hf_model_name"]
-        gguf_path = f"{output_model}/unsloth.Q4_K_M.gguf"
-
-        # Check if the GGUF file exists. If not, generate it ## TODO Multiple Quantisation other than Q4_K_M.gguf
-        if not os.path.exists(gguf_path):
-            self.model, self.tokenizer = self.load_model()
-            self.save_model_gguf()
-        return f"""FROM {output_model}/unsloth.Q4_K_M.gguf
-
-TEMPLATE \"\"\"Below are some instructions that describe some tasks. Write responses that appropriately complete each request.{{{{ if .Prompt }}}}
-
-### Instruction:
-{{{{ .Prompt }}}}
-
-{{{{ end }}}}### Response:
-{{{{ .Response }}}}\"\"\"
-
-PARAMETER stop ""
-PARAMETER stop ""
-PARAMETER stop ""
-PARAMETER stop ""
-PARAMETER stop "<|reserved_special_token_"
-"""
+        model_name = self.config["model_name"].lower()
+        # Mapping from model name keywords to their default TEMPLATE and stop tokens (and optional SYSTEM/num_ctx)
+        mapping = {
+            "llama": {
+                "template": """<|start_header_id|>system<|end_header_id|>
+    Cutting Knowledge Date: December 2023
+    {{ if .System }}{{ .System }}
+    {{- end }}
+    {{- if .Tools }}When you receive a tool call response, use the output to format an answer to the orginal user question.
+    You are a helpful assistant with tool calling capabilities.
+    {{- end }}<|eot_id|>
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 }}
+    {{- if eq .Role "user" }}<|start_header_id|>user<|end_header_id|>
+    {{- if and $.Tools $last }}
+    Given the following functions, please respond with a JSON for a function call with its proper arguments that best answers the given prompt.
+    Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. Do not use variables.
+    {{ range $.Tools }}
+    {{- . }}
+    {{ end }}
+    {{ .Content }}<|eot_id|>
+    {{- else }}
+    {{ .Content }}<|eot_id|>
+    {{- end }}{{ if $last }}<|start_header_id|>assistant<|end_header_id|>
+    {{ end }}
+    {{- else if eq .Role "assistant" }}<|start_header_id|>assistant<|end_header_id|>
+    {{- if .ToolCalls }}
+    {{ range .ToolCalls }}
+    {"name": "{{ .Function.Name }}", "parameters": {{ .Function.Arguments }}}{{ end }}
+    {{- else }}
+    {{ .Content }}
+    {{- end }}{{ if not $last }}<|eot_id|>{{ end }}
+    {{- else if eq .Role "tool" }}<|start_header_id|>ipython<|end_header_id|>
+    {{ .Content }}<|eot_id|>{{ if $last }}<|start_header_id|>assistant<|end_header_id|>
+    {{ end }}
+    {{- end }}
+    {{- end }}""",
+                "stop_tokens": ["<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>"]
+            },
+            "qwen": {
+                "template": """{{- if .Suffix }}<|fim_prefix|>{{ .Prompt }}<|fim_suffix|>{{ .Suffix }}<|fim_middle|>
+    {{- else if .Messages }}
+    {{- if or .System .Tools }}<|im_start|>system
+    {{- if .System }}
+    {{ .System }}
+    {{- end }}
+    {{- if .Tools }}
+    # Tools
+    You may call one or more functions to assist with the user query.
+    You are provided with function signatures within <tools></tools> XML tags:
+    <tools>
+    {{- range .Tools }}
+    {"type": "function", "function": {{ .Function }}}
+    {{- end }}
+    </tools>
+    For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+    <tool_call>
+    {"name": <function-name>, "arguments": <args-json-object>}
+    </tool_call>
+    {{- end }}<|im_end|>
+    {{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    {{- if eq .Role "user" }}<|im_start|>user
+    {{ .Content }}<|im_end|>
+    {{ else if eq .Role "assistant" }}<|im_start|>assistant
+    {{ if .Content }}{{ .Content }}
+    {{- else if .ToolCalls }}<tool_call>
+    {{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}
+    {{ end }}</tool_call>
+    {{- end }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- else if eq .Role "tool" }}<|im_start|>user
+    <tool_response>
+    {{ .Content }}
+    </tool_response><|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_start|>assistant
+    {{ end }}
+    {{- end }}
+    {{- else }}
+    {{- if .System }}<|im_start|>system
+    {{ .System }}<|im_end|>
+    {{ end }}{{ if .Prompt }}<|im_start|>user
+    {{ .Prompt }}<|im_end|>
+    {{ end }}<|im_start|>assistant
+    {{ end }}{{ .Response }}{{ if .Response }}<|im_end|>{{ end }}""",
+                "system": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
+                "num_ctx": 32768,
+                "stop_tokens": ["<|endoftext|>"]
+            },
+            "mistral": {
+                "template": "[INST] {{ if .System }}{{ .System }} {{ end }}{{ .Prompt }} [/INST]",
+                "stop_tokens": ["[INST]", "[/INST]"]
+            },
+            "phi": {
+                "template": """{{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    <|im_start|>{{ .Role }}<|im_sep|>
+    {{ .Content }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_end|>
+    <|im_start|>assistant<|im_sep|>
+    {{ end }}
+    {{- end }}""",
+                "stop_tokens": ["<|im_start|>", "<|im_end|>", "<|im_sep|>"]
+            },
+            "deepseek": {
+                "template": """{{- if .System }}{{ .System }}{{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1}}
+    {{- if eq .Role "user" }}
+    {{ .Content }}
+    {{- else if eq .Role "assistant" }}
+    {{ .Content }}{{- if not $last }}
+    {{- end }}
+    {{- end }}
+    {{- if and $last (ne .Role "assistant") }}
+    {{ end }}
+    {{- end }}""",
+                "stop_tokens": ["", "", "", ""]
+            },
+            "llava": {
+                "template": """{{- if .Suffix }}<|fim_prefix|>{{ .Prompt }}<|fim_suffix|>{{ .Suffix }}<|fim_middle|>
+    {{- else if .Messages }}
+    {{- if or .System .Tools }}<|im_start|>system
+    {{- if .System }}
+    {{ .System }}
+    {{- end }}
+    {{- if .Tools }}
+    # Tools
+    You may call one or more functions to assist with the user query.
+    You are provided with function signatures within <tools></tools> XML tags:
+    <tools>
+    {{- range .Tools }}
+    {"type": "function", "function": {{ .Function }}}
+    {{- end }}
+    </tools>
+    For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+    <tool_call>
+    {"name": <function-name>, "arguments": <args-json-object>}
+    </tool_call>
+    {{- end }}<|im_end|>
+    {{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    {{- if eq .Role "user" }}<|im_start|>user
+    {{ .Content }}<|im_end|>
+    {{ else if eq .Role "assistant" }}<|im_start|>assistant
+    {{ if .Content }}{{ .Content }}
+    {{- else if .ToolCalls }}<tool_call>
+    {{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}
+    {{ end }}</tool_call>
+    {{- end }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- else if eq .Role "tool" }}<|im_start|>user
+    <tool_response>
+    {{ .Content }}
+    </tool_response><|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_start|>assistant
+    {{ end }}
+    {{- end }}
+    {{- else }}
+    {{- if .System }}<|im_start|>system
+    {{ .System }}<|im_end|>
+    {{ end }}{{ if .Prompt }}<|im_start|>user
+    {{ .Prompt }}<|im_end|>
+    {{ end }}<|im_start|>assistant
+    {{ end }}{{ .Response }}{{ if .Response }}<|im_end|>{{ end }}""",
+                "stop_tokens": ["</s>", "USER:", "ASSSISTANT:"]
+            }
+        }
+        # Select mapping by checking if any key is in the model_name.
+        chosen = None
+        for key, settings in mapping.items():
+            if key in model_name:
+                chosen = settings
+                break
+        if chosen is None:
+            # Fallback default
+            chosen = {
+                "template": """{{ if .System }}<|start_header_id|>system<|end_header_id|>
+    {{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|>
+    {{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|>
+    {{ .Response }}<|eot_id|>""",
+                "stop_tokens": ["<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>"]
+            }
+        # Build the stop parameter lines.
+        stop_params = "\n".join([f"PARAMETER stop {token}" for token in chosen["stop_tokens"]])
+        # Optionally include a SYSTEM line and num_ctx if defined in the mapping.
+        system_line = ""
+        if "system" in chosen:
+            system_line = f"SYSTEM {chosen['system']}\n"
+        num_ctx_line = ""
+        if "num_ctx" in chosen:
+            num_ctx_line = f"PARAMETER num_ctx {chosen['num_ctx']}\n"
+        # Assemble and return the modelfile content.
+        return f"""FROM {output_model}
+    TEMPLATE \"\"\"{chosen['template']}\"\"\"
+    {system_line}{num_ctx_line}{stop_params}
+    """

    def create_and_push_ollama_model(self):
        modelfile_content = self.prepare_modelfile_content()
-        with open('Modelfile', 'w') as file:
+        with open("Modelfile", "w") as file:
            file.write(modelfile_content)
-
        subprocess.run(["ollama", "serve"])
        subprocess.run(["ollama", "create", f"{self.config['ollama_model']}:{self.config['model_parameters']}", "-f", "Modelfile"])
        subprocess.run(["ollama", "push", f"{self.config['ollama_model']}:{self.config['model_parameters']}"])
@@ -235,42 +533,30 @@ PARAMETER stop "<|reserved_special_token_"
        self.print_system_info()
        self.check_gpu()
        self.check_ram()
-        # self.install_packages()
        if self.config.get("train", "true").lower() == "true":
            self.prepare_model()
            self.train_model()
-
        if self.config.get("huggingface_save", "true").lower() == "true":
-            # self.model, self.tokenizer = self.load_model()
            self.save_model_merged()
-
        if self.config.get("huggingface_save_gguf", "true").lower() == "true":
-            # self.model, self.tokenizer = self.load_model()
            self.push_model_gguf()
-            
-        # if self.config.get("save_gguf", "true").lower() == "true": ## TODO
-        #     self.model, self.tokenizer = self.load_model()
-        #     self.save_model_gguf()
-        
-        # if self.config.get("save_merged", "true").lower() == "true": ## TODO
-        #     self.model, self.tokenizer = self.load_model()
-        #     self.save_model_merged()
-
        if self.config.get("ollama_save", "true").lower() == "true":
            self.create_and_push_ollama_model()

-
 def main():
    import argparse
-    parser = argparse.ArgumentParser(description='PraisonAI Training Script')
-    parser.add_argument('command', choices=['train'], help='Command to execute')
-    parser.add_argument('--config', default='config.yaml', help='Path to configuration file')
+    parser = argparse.ArgumentParser(description="PraisonAI Training Script")
+    parser.add_argument("command", choices=["train"], help="Command to execute")
+    parser.add_argument("--config", default="config.yaml", help="Path to configuration file")
+    parser.add_argument("--model", type=str, help="Model name")
+    parser.add_argument("--hf", type=str, help="Hugging Face model name")
+    parser.add_argument("--ollama", type=str, help="Ollama model name")
+    parser.add_argument("--dataset", type=str, help="Dataset name for training")
    args = parser.parse_args()

-    if args.command == 'train':
-        ai = train(config_path=args.config)
-        ai.run()
+    if args.command == "train":
+        trainer_obj = TrainModel(config_path=args.config)
+        trainer_obj.run()

-
-if __name__ == '__main__':
+if __name__ == "__main__":
    main()
@@ -0,0 +1,306 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+This script finetunes a vision language model using Unsloth's fast training framework.
+It supports vision tasks by converting raw image-caption samples into a conversation format, 
+adding vision-specific LoRA adapters, and training using TRL's SFTTrainer with UnslothVisionDataCollator.
+"""
+
+import os
+import sys
+import yaml
+import torch
+import shutil
+import subprocess
+import gc  # For garbage collection
+
+from datasets import load_dataset, concatenate_datasets, Dataset
+from unsloth import FastVisionModel, is_bf16_supported
+from unsloth.trainer import UnslothVisionDataCollator
+from transformers import TrainingArguments
+from trl import SFTTrainer
+from tqdm import tqdm  # Add progress bar
+
+
+class TrainVisionModel:
+    def __init__(self, config_path="config.yaml"):
+        self.load_config(config_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.hf_tokenizer = None  # The underlying tokenizer
+
+    def load_config(self, path):
+        with open(path, "r") as file:
+            self.config = yaml.safe_load(file)
+        print("DEBUG: Loaded config:", self.config)
+
+    def print_system_info(self):
+        print("DEBUG: PyTorch version:", torch.__version__)
+        print("DEBUG: CUDA version:", torch.version.cuda)
+        if torch.cuda.is_available():
+            print("DEBUG: CUDA Device Capability:", torch.cuda.get_device_capability())
+        else:
+            print("DEBUG: CUDA is not available")
+        print("DEBUG: Python Version:", sys.version)
+        print("DEBUG: Python Path:", sys.executable)
+
+    def check_gpu(self):
+        gpu_stats = torch.cuda.get_device_properties(0)
+        print(f"DEBUG: GPU = {gpu_stats.name}. Max memory = {round(gpu_stats.total_memory/(1024**3),3)} GB.")
+
+    def check_ram(self):
+        from psutil import virtual_memory
+        ram_gb = virtual_memory().total / 1e9
+        print(f"DEBUG: Your runtime has {ram_gb:.1f} gigabytes of available RAM")
+        if ram_gb < 20:
+            print("DEBUG: Not using a high-RAM runtime")
+        else:
+            print("DEBUG: You are using a high-RAM runtime!")
+
+    def prepare_model(self):
+        print("DEBUG: Preparing vision model and tokenizer...")
+        self.model, original_tokenizer = FastVisionModel.from_pretrained(
+            model_name=self.config["model_name"],
+            load_in_4bit=self.config["load_in_4bit"],
+            use_gradient_checkpointing="unsloth"
+        )
+        print("DEBUG: Vision model and original tokenizer loaded.")
+        
+        # Use the full processor that supports image inputs.
+        self.hf_tokenizer = original_tokenizer
+        
+        # Set pad token if needed
+        if not hasattr(self.hf_tokenizer, 'pad_token') or self.hf_tokenizer.pad_token is None:
+            if hasattr(self.hf_tokenizer, 'eos_token'):
+                self.hf_tokenizer.pad_token = self.hf_tokenizer.eos_token
+            elif hasattr(self.hf_tokenizer, 'bos_token'):
+                self.hf_tokenizer.pad_token = self.hf_tokenizer.bos_token
+        
+        # Set max length
+        if hasattr(self.hf_tokenizer, 'model_max_length'):
+            self.hf_tokenizer.model_max_length = self.config.get("max_seq_length", 2048)
+        
+        # Add vision-specific LoRA adapters
+        self.model = FastVisionModel.get_peft_model(
+            self.model,
+            finetune_vision_layers=self.config.get("finetune_vision_layers", False),
+            finetune_language_layers=self.config.get("finetune_language_layers", True),
+            finetune_attention_modules=self.config.get("finetune_attention_modules", True),
+            finetune_mlp_modules=self.config.get("finetune_mlp_modules", True),
+            r=16,
+            lora_alpha=16,
+            lora_dropout=0,
+            bias="none",
+            random_state=3407,
+            use_rslora=False,
+            loftq_config=None
+        )
+        print("DEBUG: Vision LoRA adapters added.")
+
+    def convert_sample(self, sample):
+        
+        instruction = self.config.get(
+            "vision_instruction",
+            "You are an expert radiographer. Describe accurately what you see in this image."
+        )
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": instruction},
+                    {"type": "image", "image": sample["image"]}
+                ]
+            },
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": sample["caption"]}
+                ]
+            },
+        ]
+        
+        return {"messages": conversation}
+
+    def load_datasets(self):
+        all_converted = []
+        for dataset_info in self.config["dataset"]:
+            print("\nDEBUG: Loading vision dataset:", dataset_info)
+            ds = load_dataset(
+                dataset_info["name"], 
+                split=dataset_info.get("split_type", "train")
+            )
+            print("DEBUG: Dataset size:", len(ds))
+            print("DEBUG: First raw sample:", ds[0])
+            print("DEBUG: Dataset features:", ds.features)
+            
+            print("\nDEBUG: Converting dataset to vision conversation format...")
+            converted_ds = [self.convert_sample(sample) for sample in ds]
+            
+            # Debug first converted sample
+            print("\nDEBUG: First converted sample structure:")
+            first = converted_ds[0]
+            print("DEBUG: Message keys:", first["messages"][0]["content"][1].keys())
+            print("DEBUG: Image type in converted:", type(first["messages"][0]["content"][1].get("image")))
+            
+            all_converted.extend(converted_ds)
+        
+        print("\nDEBUG: Combined vision dataset has", len(all_converted), "examples.")
+        return all_converted
+
+    def train_model(self):
+        print("DEBUG: Starting vision training...")
+        raw_dataset = self.load_datasets()
+        
+        # Build training arguments using TrainingArguments
+        training_args = TrainingArguments(
+            per_device_train_batch_size=self.config.get("per_device_train_batch_size", 1),
+            gradient_accumulation_steps=self.config.get("gradient_accumulation_steps", 4),
+            warmup_steps=self.config.get("warmup_steps", 5),
+            max_steps=self.config.get("max_steps", 30),
+            learning_rate=self.config.get("learning_rate", 2e-4),
+            fp16=self.config.get("fp16", not is_bf16_supported()),
+            bf16=self.config.get("bf16", is_bf16_supported()),
+            logging_steps=self.config.get("logging_steps", 1),
+            optim=self.config.get("optim", "adamw_8bit"),
+            weight_decay=self.config.get("weight_decay", 0.01),
+            lr_scheduler_type=self.config.get("lr_scheduler_type", "linear"),
+            seed=self.config.get("seed", 3407),
+            output_dir=self.config.get("output_dir", "outputs"),
+            report_to="none" if not os.getenv("PRAISON_WANDB") else "wandb",
+            remove_unused_columns=False,
+            # Add memory optimization settings
+            gradient_checkpointing=True,
+            max_grad_norm=1.0,
+        )
+        
+        trainer = SFTTrainer(
+            model=self.model,
+            tokenizer=self.hf_tokenizer,
+            data_collator=UnslothVisionDataCollator(self.model, self.hf_tokenizer),
+            train_dataset=raw_dataset,
+            args=training_args,
+            max_seq_length=self.config.get("max_seq_length", 2048),
+            dataset_text_field="",  # Required for vision training
+            dataset_kwargs={"skip_prepare_dataset": True},  # Required for vision training
+            packing=False  # Explicitly set packing to False
+        )
+        print("DEBUG: Beginning vision trainer.train() ...")
+        trainer.train()
+        print("DEBUG: Vision training complete. Saving model and tokenizer locally...")
+        self.model.save_pretrained("lora_vision_model")
+        self.hf_tokenizer.save_pretrained("lora_vision_model")
+        print("DEBUG: Saved vision model and tokenizer to 'lora_vision_model'.")
+
+    def vision_inference(self, instruction, image):
+        FastVisionModel.for_inference(self.model)
+        messages = [
+            {"role": "user", "content": [
+                {"type": "image"},
+                {"type": "text", "text": instruction}
+            ]}
+        ]
+        input_text = self.hf_tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+        inputs = self.hf_tokenizer(
+            image,
+            input_text,
+            add_special_tokens=False,
+            return_tensors="pt"
+        ).to("cuda")
+        outputs = self.model.generate(
+            **inputs,
+            max_new_tokens=128,
+            use_cache=True,
+            temperature=1.5,
+            min_p=0.1
+        )
+        print("DEBUG: Vision inference output:", self.hf_tokenizer.batch_decode(outputs))
+
+    def save_model_merged(self):
+        if os.path.exists(self.config["hf_model_name"]):
+            shutil.rmtree(self.config["hf_model_name"])
+        self.model.push_to_hub_merged(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            save_method="merged_16bit",
+            token=os.getenv("HF_TOKEN")
+        )
+
+    def push_model_gguf(self):
+        self.model.push_to_hub_gguf(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            quantization_method=self.config.get("quantization_method", "q4_k_m"),
+            token=os.getenv("HF_TOKEN")
+        )
+
+    def save_model_gguf(self):
+        self.model.save_pretrained_gguf(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            quantization_method="q4_k_m"
+        )
+
+    def prepare_modelfile_content(self):
+        output_model = self.config["hf_model_name"]
+        
+        template = '''{{- range $index, $_ := .Messages }}<|start_header_id|>{{ .Role }}<|end_header_id|>
+
+{{ .Content }}
+{{- if gt (len (slice $.Messages $index)) 1 }}<|eot_id|>
+{{- else if ne .Role "assistant" }}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+{{ end }}
+{{- end }}'''
+        
+        return f"""FROM {output_model}
+TEMPLATE {template}
+PARAMETER temperature 0.6
+PARAMETER top_p 0.9
+"""
+
+    def create_and_push_ollama_model(self):
+        modelfile_content = self.prepare_modelfile_content()
+        with open("Modelfile", "w") as file:
+            file.write(modelfile_content)
+        subprocess.run(["ollama", "serve"])
+        subprocess.run(["ollama", "create", f"{self.config['ollama_model']}:{self.config['model_parameters']}", "-f", "Modelfile"])
+        subprocess.run(["ollama", "push", f"{self.config['ollama_model']}:{self.config['model_parameters']}"])
+
+    def run(self):
+        self.print_system_info()
+        self.check_gpu()
+        self.check_ram()
+        if self.config.get("train", "true").lower() == "true":
+            self.prepare_model()
+            self.train_model()
+        if self.config.get("huggingface_save", "true").lower() == "true":
+            self.save_model_merged()
+        if self.config.get("huggingface_save_gguf", "true").lower() == "true":
+            self.push_model_gguf()
+        if self.config.get("ollama_save", "true").lower() == "true":
+            self.create_and_push_ollama_model()
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="PraisonAI Vision Training Script")
+    parser.add_argument("command", choices=["train", "inference"], help="Command to execute")
+    parser.add_argument("--config", default="config.yaml", help="Path to configuration file")
+    args = parser.parse_args()
+
+    trainer_obj = TrainVisionModel(config_path=args.config)
+    if args.command == "train":
+        trainer_obj.run()
+    elif args.command == "inference":
+        # For inference, we load a sample image from the first dataset
+        instr = trainer_obj.config.get("vision_instruction", "You are an expert radiographer. Describe accurately what you see in this image.")
+        ds_info = trainer_obj.config["dataset"][0]
+        ds = load_dataset(ds_info["name"], split=ds_info.get("split_type", "train"))
+        sample_image = ds[0]["image"]
+        if trainer_obj.model is None or trainer_obj.hf_tokenizer is None:
+            trainer_obj.prepare_model()
+        trainer_obj.vision_inference(instr, sample_image)
+
+
+if __name__ == "__main__":
+    main() 
@@ -0,0 +1,140 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+This script handles uploading trained vision models to Hugging Face and Ollama.
+It reads configuration from config.yaml and provides options to upload in different formats.
+"""
+
+import os
+import yaml
+import torch
+import shutil
+import subprocess
+from unsloth import FastVisionModel
+
+class UploadVisionModel:
+    def __init__(self, config_path="config.yaml"):
+        self.load_config(config_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.hf_tokenizer = None
+
+    def load_config(self, path):
+        """Load configuration from yaml file."""
+        with open(path, "r") as file:
+            self.config = yaml.safe_load(file)
+        print("DEBUG: Loaded config:", self.config)
+
+    def prepare_model(self):
+        """Load the trained model for uploading."""
+        print("DEBUG: Loading trained model and tokenizer...")
+        self.model, original_tokenizer = FastVisionModel.from_pretrained(
+            model_name=self.config.get("output_dir", "lora_model"),
+            load_in_4bit=self.config.get("load_in_4bit", True)
+        )
+        self.hf_tokenizer = original_tokenizer
+        print("DEBUG: Model and tokenizer loaded successfully.")
+
+    def save_model_merged(self):
+        """Save merged model to Hugging Face Hub."""
+        print(f"DEBUG: Saving merged model to Hugging Face Hub: {self.config['hf_model_name']}")
+        if os.path.exists(self.config["hf_model_name"]):
+            shutil.rmtree(self.config["hf_model_name"])
+        self.model.push_to_hub_merged(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            save_method="merged_16bit",
+            token=os.getenv("HF_TOKEN")
+        )
+        print("DEBUG: Model saved to Hugging Face Hub successfully.")
+
+    def push_model_gguf(self):
+        """Push model in GGUF format to Hugging Face Hub."""
+        print(f"DEBUG: Pushing GGUF model to Hugging Face Hub: {self.config['hf_model_name']}")
+        self.model.push_to_hub_gguf(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            quantization_method=self.config.get("quantization_method", "q4_k_m"),
+            token=os.getenv("HF_TOKEN")
+        )
+        print("DEBUG: GGUF model pushed to Hugging Face Hub successfully.")
+
+    def prepare_modelfile_content(self):
+        """Prepare Ollama modelfile content using Llama 3.2 vision template."""
+        output_model = self.config["hf_model_name"]
+        
+        # Using Llama 3.2 vision template format
+        template = """{{- range $index, $_ := .Messages }}<|start_header_id|>{{ .Role }}<|end_header_id|>
+
+{{ .Content }}
+{{- if gt (len (slice $.Messages $index)) 1 }}<|eot_id|>
+{{- else if ne .Role "assistant" }}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+{{ end }}
+{{- end }}"""
+        
+        # Assemble the modelfile content with Llama 3.2 vision parameters
+        modelfile = f"FROM {output_model}\n"
+        modelfile += "TEMPLATE \"""" + template + "\"""\n"
+        modelfile += "PARAMETER temperature 0.6\n"
+        modelfile += "PARAMETER top_p 0.9\n"
+        return modelfile
+
+    def create_and_push_ollama_model(self):
+        """Create and push model to Ollama."""
+        print(f"DEBUG: Creating Ollama model: {self.config['ollama_model']}:{self.config['model_parameters']}")
+        modelfile_content = self.prepare_modelfile_content()
+        with open("Modelfile", "w") as file:
+            file.write(modelfile_content)
+        
+        print("DEBUG: Starting Ollama server...")
+        subprocess.run(["ollama", "serve"])
+        
+        print("DEBUG: Creating Ollama model...")
+        subprocess.run([
+            "ollama", "create", 
+            f"{self.config['ollama_model']}:{self.config['model_parameters']}", 
+            "-f", "Modelfile"
+        ])
+        
+        print("DEBUG: Pushing model to Ollama...")
+        subprocess.run([
+            "ollama", "push", 
+            f"{self.config['ollama_model']}:{self.config['model_parameters']}"
+        ])
+        print("DEBUG: Model pushed to Ollama successfully.")
+
+    def upload(self, target="all"):
+        """
+        Upload the model to specified targets.
+        Args:
+            target (str): One of 'all', 'huggingface', 'huggingface_gguf', or 'ollama'
+        """
+        self.prepare_model()
+        
+        if target in ["all", "huggingface"]:
+            self.save_model_merged()
+        
+        if target in ["all", "huggingface_gguf"]:
+            self.push_model_gguf()
+            
+        if target in ["all", "ollama"]:
+            self.create_and_push_ollama_model()
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="Upload Vision Model to Various Platforms")
+    parser.add_argument("--config", default="config.yaml", help="Path to configuration file")
+    parser.add_argument(
+        "--target", 
+        choices=["all", "huggingface", "huggingface_gguf", "ollama"],
+        default="all",
+        help="Target platform to upload to"
+    )
+    args = parser.parse_args()
+
+    uploader = UploadVisionModel(config_path=args.config)
+    uploader.upload(target=args.target)
+
+if __name__ == "__main__":
+    main()
@@ -1,6 +1,6 @@
 [project]
 name = "PraisonAI"
-version = "2.0.62"
+version = "2.0.80"
 description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human-agent collaboration."
 readme = "README.md"
 license = ""
@@ -12,7 +12,7 @@ dependencies = [
    "rich>=13.7",
    "markdown>=3.5",
    "pyparsing>=3.0.0",
-    "praisonaiagents>=0.0.61",
+    "praisonaiagents>=0.0.65",
    "python-dotenv>=0.19.0",
    "instructor>=1.3.3",
    "PyYAML>=6.0",
@@ -84,7 +84,7 @@ autogen = ["pyautogen>=0.2.19", "praisonai-tools>=0.0.7", "crewai"]

 [tool.poetry]
 name = "PraisonAI"
-version = "2.0.62"
+version = "2.0.80"
 description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human–agent collaboration."
 authors = ["Mervin Praison"]
 license = ""
@@ -102,7 +102,7 @@ python = ">=3.10,<3.13"
 rich = ">=13.7"
 markdown = ">=3.5"
 pyparsing = ">=3.0.0"
-praisonaiagents = ">=0.0.61"
+praisonaiagents = ">=0.0.65"
 python-dotenv = ">=0.19.0"
 instructor = ">=1.3.3"
 PyYAML = ">=6.0"
@@ -0,0 +1,37 @@
+from praisonaiagents import Agent
+
+agent = Agent(
+    instructions="You are a Wikipedia Agent", 
+    llm="anthropic/claude-3-7-sonnet-20250219"
+)
+result = agent.start("Why Sky is Blue?")
+print(result)
+
+from praisonaiagents import Agent
+from praisonaiagents.tools import internet_search
+
+agent = Agent(
+    instructions="You are a Wikipedia Agent", 
+    tools=[internet_search],
+    llm="anthropic/claude-3-7-sonnet-20250219"
+)
+agent.start("What is Praison AI?")
+
+from praisonaiagents import Agent, PraisonAIAgents
+from praisonaiagents.tools import internet_search
+
+
+research_agent = Agent(
+    instructions="Search Information about Claude Sonnet 3.7", 
+    tools=[internet_search],
+    llm="anthropic/claude-3-7-sonnet-20250219"
+)
+
+editor_agent = Agent(
+    instructions="Write a Blog Post with the provided information about Claude Sonnet 3.7", 
+    llm="anthropic/claude-3-7-sonnet-20250219"
+)
+
+agents = PraisonAIAgents(agents=[research_agent, editor_agent])
+result = agents.start()
+print(result)
@@ -1,10 +1,10 @@
 from praisonaiagents import Agent
+from praisonaiagents.tools import internet_search

 agent = Agent(
    instructions="You are a helpful assistant",
-    llm="gemini/gemini-1.5-flash-8b",
-    self_reflect=True,
-    verbose=True
+    tools=[internet_search],
+    llm="gemini/gemini-1.5-flash-8b"
 )

-agent.start("Why sky is Blue?")
+agent.start("What is Praison AI?")
@@ -0,0 +1,18 @@
+from praisonaiagents import Agent
+from praisonaiagents.tools import wiki_search, wiki_summary, wiki_page, wiki_random, wiki_language
+
+agent1 = Agent(
+    instructions="You are a Wikipedia Agent", 
+    tools=[wiki_search, wiki_summary, wiki_page, wiki_random, wiki_language],
+    llm="openai/gpt-4o-mini",
+    verbose=10
+)
+agent1.start("history of AI in 1 line")
+
+agent2 = Agent(
+    instructions="You are a Wikipedia Agent", 
+    tools=[wiki_search, wiki_summary, wiki_page, wiki_random, wiki_language],
+    llm="gpt-4o-mini",
+    verbose=10
+)
+agent2.start("history of AI in 1 line")
@@ -714,6 +714,22 @@ Your Goal: {self.goal}
            return None

    def chat(self, prompt, temperature=0.2, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False):
+        # Log all parameter values when in debug mode
+        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+            param_info = {
+                "prompt": str(prompt)[:100] + "..." if isinstance(prompt, str) and len(str(prompt)) > 100 else str(prompt),
+                "temperature": temperature,
+                "tools": [t.__name__ if hasattr(t, "__name__") else str(t) for t in tools] if tools else None,
+                "output_json": str(output_json.__class__.__name__) if output_json else None,
+                "output_pydantic": str(output_pydantic.__class__.__name__) if output_pydantic else None,
+                "reasoning_steps": reasoning_steps,
+                "agent_name": self.name,
+                "agent_role": self.role,
+                "agent_goal": self.goal
+            }
+            logging.debug(f"Agent.chat parameters: {json.dumps(param_info, indent=2, default=str)}")
+        
+        start_time = time.time()
        reasoning_steps = reasoning_steps or self.reasoning_steps
        # Search for existing knowledge if any knowledge is provided
        if self.knowledge:
@@ -738,7 +754,7 @@ Your Goal: {self.goal}
                    system_prompt=f"{self.backstory}\n\nYour Role: {self.role}\n\nYour Goal: {self.goal}" if self.use_system_prompt else None,
                    chat_history=self.chat_history,
                    temperature=temperature,
-                    tools=tools,
+                    tools=self.tools if tools is None else tools,
                    output_json=output_json,
                    output_pydantic=output_pydantic,
                    verbose=self.verbose,
@@ -749,7 +765,7 @@ Your Goal: {self.goal}
                    console=self.console,
                    agent_name=self.name,
                    agent_role=self.role,
-                    agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in self.tools],
+                    agent_tools=[t.__name__ if hasattr(t, '__name__') else str(t) for t in (tools if tools is not None else self.tools)],
                    execute_tool_fn=self.execute_tool,  # Pass tool execution function
                    reasoning_steps=reasoning_steps
                )
@@ -757,6 +773,11 @@ Your Goal: {self.goal}
                self.chat_history.append({"role": "user", "content": prompt})
                self.chat_history.append({"role": "assistant", "content": response_text})

+                # Log completion time if in debug mode
+                if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                    total_time = time.time() - start_time
+                    logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")
+
                return response_text
            except Exception as e:
                display_error(f"Error in LLM chat: {e}")
@@ -944,6 +965,13 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                    display_error(f"Error in chat: {e}", console=self.console)
                    return None 

+        # Log completion time if in debug mode
+        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+            total_time = time.time() - start_time
+            logging.debug(f"Agent.chat completed in {total_time:.2f} seconds")
+            
+        return response_text
+
    def clean_json_output(self, output: str) -> str:
        """Clean and extract JSON from response text."""
        cleaned = output.strip()
@@ -958,6 +986,22 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.

    async def achat(self, prompt: str, temperature=0.2, tools=None, output_json=None, output_pydantic=None, reasoning_steps=False):
        """Async version of chat method. TODO: Requires Syncing with chat method.""" 
+        # Log all parameter values when in debug mode
+        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+            param_info = {
+                "prompt": str(prompt)[:100] + "..." if isinstance(prompt, str) and len(str(prompt)) > 100 else str(prompt),
+                "temperature": temperature,
+                "tools": [t.__name__ if hasattr(t, "__name__") else str(t) for t in tools] if tools else None,
+                "output_json": str(output_json.__class__.__name__) if output_json else None,
+                "output_pydantic": str(output_pydantic.__class__.__name__) if output_pydantic else None,
+                "reasoning_steps": reasoning_steps,
+                "agent_name": self.name,
+                "agent_role": self.role,
+                "agent_goal": self.goal
+            }
+            logging.debug(f"Agent.achat parameters: {json.dumps(param_info, indent=2, default=str)}")
+        
+        start_time = time.time()
        reasoning_steps = reasoning_steps or self.reasoning_steps
        try:
            # Search for existing knowledge if any knowledge is provided
@@ -996,9 +1040,15 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                    self.chat_history.append({"role": "user", "content": prompt})
                    self.chat_history.append({"role": "assistant", "content": response_text})

+                    if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                        total_time = time.time() - start_time
+                        logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                    return response_text
                except Exception as e:
                    display_error(f"Error in LLM chat: {e}")
+                    if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                        total_time = time.time() - start_time
+                        logging.debug(f"Agent.achat failed in {total_time:.2f} seconds: {str(e)}")
                    return None

            # For OpenAI client
@@ -1081,7 +1131,11 @@ Your Goal: {self.goal}
                            temperature=temperature,
                            tools=formatted_tools
                        )
-                        return await self._achat_completion(response, tools)
+                        result = await self._achat_completion(response, tools)
+                        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                            total_time = time.time() - start_time
+                            logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
+                        return result
                    elif output_json or output_pydantic:
                        response = await async_client.chat.completions.create(
                            model=self.llm,
@@ -1090,6 +1144,9 @@ Your Goal: {self.goal}
                            response_format={"type": "json_object"}
                        )
                        # Return the raw response
+                        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                            total_time = time.time() - start_time
+                            logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                        return response.choices[0].message.content
                    else:
                        response = await async_client.chat.completions.create(
@@ -1097,12 +1154,21 @@ Your Goal: {self.goal}
                            messages=messages,
                            temperature=temperature
                        )
+                        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                            total_time = time.time() - start_time
+                            logging.debug(f"Agent.achat completed in {total_time:.2f} seconds")
                        return response.choices[0].message.content
                except Exception as e:
                    display_error(f"Error in chat completion: {e}")
+                    if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                        total_time = time.time() - start_time
+                        logging.debug(f"Agent.achat failed in {total_time:.2f} seconds: {str(e)}")
                    return None
        except Exception as e:
            display_error(f"Error in achat: {e}")
+            if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                total_time = time.time() - start_time
+                logging.debug(f"Agent.achat failed in {total_time:.2f} seconds: {str(e)}")
            return None

    async def _achat_completion(self, response, tools, reasoning_steps=False):
@@ -477,8 +477,14 @@ Context:
                else:
                    self.run_task(task_id)

-    async def astart(self, content=None, **kwargs):
-        """Async version of start method"""
+    async def astart(self, content=None, return_dict=False, **kwargs):
+        """Async version of start method
+        
+        Args:
+            content: Optional content to add to all tasks' context
+            return_dict: If True, returns the full results dictionary instead of only the final response
+            **kwargs: Additional arguments
+        """
        if content:
            # Add content to context of all tasks
            for task in self.tasks.values():
@@ -488,10 +494,25 @@ Context:
                    task.context.append(content)

        await self.arun_all_tasks()
-        return {
+        
+        # Get results
+        results = {
            "task_status": self.get_all_tasks_status(),
            "task_results": {task_id: self.get_task_result(task_id) for task_id in self.tasks}
        }
+        
+        # By default, return only the final agent's response
+        if not return_dict:
+            # Get the last task (assuming sequential processing)
+            task_ids = list(self.tasks.keys())
+            if task_ids:
+                last_task_id = task_ids[-1]
+                last_result = self.get_task_result(last_task_id)
+                if last_result:
+                    return last_result.raw
+                    
+        # Return full results dict if return_dict is True or if no final result was found
+        return results

    def save_output_to_file(self, task, task_output):
        if task.output_file:
@@ -801,8 +822,14 @@ Context:
            return str(agent[0])
        return None

-    def start(self, content=None, **kwargs):
-        """Start agent execution with optional content and config"""
+    def start(self, content=None, return_dict=False, **kwargs):
+        """Start agent execution with optional content and config
+        
+        Args:
+            content: Optional content to add to all tasks' context
+            return_dict: If True, returns the full results dictionary instead of only the final response
+            **kwargs: Additional arguments
+        """
        if content:
            # Add content to context of all tasks
            for task in self.tasks.values():
@@ -815,10 +842,25 @@ Context:
                
        # Run tasks as before
        self.run_all_tasks()
-        return {
+        
+        # Get results
+        results = {
            "task_status": self.get_all_tasks_status(),
            "task_results": {task_id: self.get_task_result(task_id) for task_id in self.tasks}
        }
+        
+        # By default, return only the final agent's response
+        if not return_dict:
+            # Get the last task (assuming sequential processing)
+            task_ids = list(self.tasks.keys())
+            if task_ids:
+                last_task_id = task_ids[-1]
+                last_result = self.get_task_result(last_task_id)
+                if last_result:
+                    return last_result.raw
+                    
+        # Return full results dict if return_dict is True or if no final result was found
+        return results

    def set_state(self, key: str, value: Any) -> None:
        """Set a state value"""
@@ -171,7 +171,39 @@ class LLM:
        
        # Enable error dropping for cleaner output
        litellm.drop_params = True
+        # Enable parameter modification for providers like Anthropic
+        litellm.modify_params = True
        self._setup_event_tracking(events)
+        
+        # Log all initialization parameters when in debug mode
+        if not isinstance(verbose, bool) and verbose >= 10:
+            debug_info = {
+                "model": self.model,
+                "timeout": self.timeout,
+                "temperature": self.temperature,
+                "top_p": self.top_p,
+                "n": self.n,
+                "max_tokens": self.max_tokens,
+                "presence_penalty": self.presence_penalty,
+                "frequency_penalty": self.frequency_penalty,
+                "logit_bias": self.logit_bias,
+                "response_format": self.response_format,
+                "seed": self.seed,
+                "logprobs": self.logprobs,
+                "top_logprobs": self.top_logprobs,
+                "api_version": self.api_version,
+                "stop_phrases": self.stop_phrases,
+                "api_key": "***" if self.api_key else None,  # Mask API key for security
+                "base_url": self.base_url,
+                "verbose": self.verbose,
+                "markdown": self.markdown,
+                "self_reflect": self.self_reflect,
+                "max_reflect": self.max_reflect,
+                "min_reflect": self.min_reflect,
+                "reasoning_steps": self.reasoning_steps,
+                "extra_settings": {k: v for k, v in self.extra_settings.items() if k not in ["api_key"]}
+            }
+            logging.debug(f"LLM instance initialized with: {json.dumps(debug_info, indent=2, default=str)}")

    def get_response(
        self,
@@ -195,6 +227,56 @@ class LLM:
        **kwargs
    ) -> str:
        """Enhanced get_response with all OpenAI-like features"""
+        logging.info(f"Getting response from {self.model}")
+        # Log all self values when in debug mode
+        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+            debug_info = {
+                "model": self.model,
+                "timeout": self.timeout,
+                "temperature": self.temperature,
+                "top_p": self.top_p,
+                "n": self.n,
+                "max_tokens": self.max_tokens,
+                "presence_penalty": self.presence_penalty,
+                "frequency_penalty": self.frequency_penalty,
+                "logit_bias": self.logit_bias,
+                "response_format": self.response_format,
+                "seed": self.seed,
+                "logprobs": self.logprobs,
+                "top_logprobs": self.top_logprobs,
+                "api_version": self.api_version,
+                "stop_phrases": self.stop_phrases,
+                "api_key": "***" if self.api_key else None,  # Mask API key for security
+                "base_url": self.base_url,
+                "verbose": self.verbose,
+                "markdown": self.markdown,
+                "self_reflect": self.self_reflect,
+                "max_reflect": self.max_reflect,
+                "min_reflect": self.min_reflect,
+                "reasoning_steps": self.reasoning_steps
+            }
+            logging.debug(f"LLM instance configuration: {json.dumps(debug_info, indent=2, default=str)}")
+            
+            # Log the parameter values passed to get_response
+            param_info = {
+                "prompt": str(prompt)[:100] + "..." if isinstance(prompt, str) and len(str(prompt)) > 100 else str(prompt),
+                "system_prompt": system_prompt[:100] + "..." if system_prompt and len(system_prompt) > 100 else system_prompt,
+                "chat_history": f"[{len(chat_history)} messages]" if chat_history else None,
+                "temperature": temperature,
+                "tools": [t.__name__ if hasattr(t, "__name__") else str(t) for t in tools] if tools else None,
+                "output_json": str(output_json.__class__.__name__) if output_json else None,
+                "output_pydantic": str(output_pydantic.__class__.__name__) if output_pydantic else None,
+                "verbose": verbose,
+                "markdown": markdown,
+                "self_reflect": self_reflect,
+                "max_reflect": max_reflect,
+                "min_reflect": min_reflect,
+                "agent_name": agent_name,
+                "agent_role": agent_role,
+                "agent_tools": agent_tools,
+                "kwargs": str(kwargs)
+            }
+            logging.debug(f"get_response parameters: {json.dumps(param_info, indent=2, default=str)}")
        try:
            import litellm
            # This below **kwargs** is passed to .completion() directly. so reasoning_steps has to be popped. OR find alternate best way of handling this.
@@ -202,6 +284,23 @@ class LLM:
            # Disable litellm debug messages
            litellm.set_verbose = False
            
+            # Format tools if provided
+            formatted_tools = None
+            if tools:
+                formatted_tools = []
+                for tool in tools:
+                    if callable(tool):
+                        tool_def = self._generate_tool_definition(tool.__name__)
+                    elif isinstance(tool, str):
+                        tool_def = self._generate_tool_definition(tool)
+                    else:
+                        continue
+                        
+                    if tool_def:
+                        formatted_tools.append(tool_def)
+                if not formatted_tools:
+                    formatted_tools = None
+            
            # Build messages list
            messages = []
            if system_prompt:
@@ -260,6 +359,7 @@ class LLM:
                            messages=messages,
                            temperature=temperature,
                            stream=False,  # force non-streaming
+                            tools=formatted_tools,
                            **{k:v for k,v in kwargs.items() if k != 'reasoning_steps'}
                        )
                        reasoning_content = resp["choices"][0]["message"].get("provider_specific_fields", {}).get("reasoning_content")
@@ -291,6 +391,7 @@ class LLM:
                                for chunk in litellm.completion(
                                    model=self.model,
                                    messages=messages,
+                                    tools=formatted_tools,
                                    temperature=temperature,
                                    stream=True,
                                    **kwargs
@@ -305,6 +406,7 @@ class LLM:
                            for chunk in litellm.completion(
                                model=self.model,
                                messages=messages,
+                                tools=formatted_tools,
                                temperature=temperature,
                                stream=True,
                                **kwargs
@@ -318,6 +420,7 @@ class LLM:
                    final_response = litellm.completion(
                        model=self.model,
                        messages=messages,
+                        tools=formatted_tools,
                        temperature=temperature,
                        stream=False,  # No streaming for tool call check
                        **kwargs
@@ -337,20 +440,29 @@ class LLM:
                            function_name = tool_call["function"]["name"]
                            arguments = json.loads(tool_call["function"]["arguments"])

-                            if verbose:
-                                display_tool_call(f"Agent {agent_name} is calling function '{function_name}' with arguments: {arguments}", console=console)
-
+                            logging.debug(f"[TOOL_EXEC_DEBUG] About to execute tool {function_name} with args: {arguments}")
                            tool_result = execute_tool_fn(function_name, arguments)
+                            logging.debug(f"[TOOL_EXEC_DEBUG] Tool execution result: {tool_result}")

-                            if tool_result:
-                                if verbose:
-                                    display_tool_call(f"Function '{function_name}' returned: {tool_result}", console=console)
+                            if verbose:
+                                display_message = f"Agent {agent_name} called function '{function_name}' with arguments: {arguments}\n"
+                                if tool_result:
+                                    display_message += f"Function returned: {tool_result}"
+                                    logging.debug(f"[TOOL_EXEC_DEBUG] Display message with result: {display_message}")
+                                else:
+                                    display_message += "Function returned no output"
+                                    logging.debug("[TOOL_EXEC_DEBUG] Tool returned no output")
+                                
+                                logging.debug(f"[TOOL_EXEC_DEBUG] About to display tool call with message: {display_message}")
+                                display_tool_call(display_message, console=console)
+                                
                                messages.append({
                                    "role": "tool",
                                    "tool_call_id": tool_call["id"],
                                    "content": json.dumps(tool_result)
                                })
                            else:
+                                logging.debug("[TOOL_EXEC_DEBUG] Verbose mode off, not displaying tool call")
                                messages.append({
                                    "role": "tool",
                                    "tool_call_id": tool_call["id"],
@@ -552,6 +664,11 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
        except Exception as error:
            display_error(f"Error in get_response: {str(error)}")
            raise
+        
+        # Log completion time if in debug mode
+        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+            total_time = time.time() - start_time
+            logging.debug(f"get_response completed in {total_time:.2f} seconds")

    async def get_response_async(
        self,
@@ -577,6 +694,56 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
        """Async version of get_response with identical functionality."""
        try:
            import litellm
+            logging.info(f"Getting async response from {self.model}")
+            # Log all self values when in debug mode
+            if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                debug_info = {
+                    "model": self.model,
+                    "timeout": self.timeout,
+                    "temperature": self.temperature,
+                    "top_p": self.top_p,
+                    "n": self.n,
+                    "max_tokens": self.max_tokens,
+                    "presence_penalty": self.presence_penalty,
+                    "frequency_penalty": self.frequency_penalty,
+                    "logit_bias": self.logit_bias,
+                    "response_format": self.response_format,
+                    "seed": self.seed,
+                    "logprobs": self.logprobs,
+                    "top_logprobs": self.top_logprobs,
+                    "api_version": self.api_version,
+                    "stop_phrases": self.stop_phrases,
+                    "api_key": "***" if self.api_key else None,  # Mask API key for security
+                    "base_url": self.base_url,
+                    "verbose": self.verbose,
+                    "markdown": self.markdown,
+                    "self_reflect": self.self_reflect,
+                    "max_reflect": self.max_reflect,
+                    "min_reflect": self.min_reflect,
+                    "reasoning_steps": self.reasoning_steps
+                }
+                logging.debug(f"LLM async instance configuration: {json.dumps(debug_info, indent=2, default=str)}")
+                
+                # Log the parameter values passed to get_response_async
+                param_info = {
+                    "prompt": str(prompt)[:100] + "..." if isinstance(prompt, str) and len(str(prompt)) > 100 else str(prompt),
+                    "system_prompt": system_prompt[:100] + "..." if system_prompt and len(system_prompt) > 100 else system_prompt,
+                    "chat_history": f"[{len(chat_history)} messages]" if chat_history else None,
+                    "temperature": temperature,
+                    "tools": [t.__name__ if hasattr(t, "__name__") else str(t) for t in tools] if tools else None,
+                    "output_json": str(output_json.__class__.__name__) if output_json else None,
+                    "output_pydantic": str(output_pydantic.__class__.__name__) if output_pydantic else None,
+                    "verbose": verbose,
+                    "markdown": markdown,
+                    "self_reflect": self_reflect,
+                    "max_reflect": max_reflect,
+                    "min_reflect": min_reflect,
+                    "agent_name": agent_name,
+                    "agent_role": agent_role,
+                    "agent_tools": agent_tools,
+                    "kwargs": str(kwargs)
+                }
+                logging.debug(f"get_response_async parameters: {json.dumps(param_info, indent=2, default=str)}")
            reasoning_steps = kwargs.pop('reasoning_steps', self.reasoning_steps)
            litellm.set_verbose = False

@@ -767,14 +934,15 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                        function_name = tool_call.function.name
                        arguments = json.loads(tool_call.function.arguments)

-                        if verbose:
-                            display_tool_call(f"Agent {agent_name} is calling function '{function_name}' with arguments: {arguments}", console=console)
-
                        tool_result = await execute_tool_fn(function_name, arguments)

-                        if tool_result:
-                            if verbose:
-                                display_tool_call(f"Function '{function_name}' returned: {tool_result}", console=console)
+                        if verbose:
+                            display_message = f"Agent {agent_name} called function '{function_name}' with arguments: {arguments}\n"
+                            if tool_result:
+                                display_message += f"Function returned: {tool_result}"
+                            else:
+                                display_message += "Function returned no output"
+                            display_tool_call(display_message, console=console)
                            messages.append({
                                "role": "tool",
                                "tool_call_id": tool_call.id,
@@ -983,6 +1151,11 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
                raise LLMContextLengthExceededException(str(error))
            display_error(f"Error in get_response_async: {str(error)}")
            raise
+            
+        # Log completion time if in debug mode
+        if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+            total_time = time.time() - start_time
+            logging.debug(f"get_response_async completed in {total_time:.2f} seconds")

    def can_use_tools(self) -> bool:
        """Check if this model can use tool functions"""
@@ -1065,6 +1238,24 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
            
            logger.debug("Using synchronous response function")
            
+            # Log all self values when in debug mode
+            if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                debug_info = {
+                    "model": self.model,
+                    "timeout": self.timeout,
+                    "temperature": temperature,
+                    "top_p": self.top_p,
+                    "n": self.n,
+                    "max_tokens": self.max_tokens,
+                    "presence_penalty": self.presence_penalty,
+                    "frequency_penalty": self.frequency_penalty,
+                    "stream": stream,
+                    "verbose": verbose,
+                    "markdown": markdown,
+                    "kwargs": str(kwargs)
+                }
+                logger.debug(f"Response method configuration: {json.dumps(debug_info, indent=2, default=str)}")
+            
            # Build messages list
            messages = []
            if system_prompt:
@@ -1150,6 +1341,24 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
            
            logger.debug("Using asynchronous response function")
            
+            # Log all self values when in debug mode
+            if logging.getLogger().getEffectiveLevel() == logging.DEBUG:
+                debug_info = {
+                    "model": self.model,
+                    "timeout": self.timeout,
+                    "temperature": temperature,
+                    "top_p": self.top_p,
+                    "n": self.n,
+                    "max_tokens": self.max_tokens,
+                    "presence_penalty": self.presence_penalty,
+                    "frequency_penalty": self.frequency_penalty,
+                    "stream": stream,
+                    "verbose": verbose,
+                    "markdown": markdown,
+                    "kwargs": str(kwargs)
+                }
+                logger.debug(f"Async response method configuration: {json.dumps(debug_info, indent=2, default=str)}")
+            
            # Build messages list
            messages = []
            if system_prompt:
@@ -1210,4 +1419,117 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.

        except Exception as error:
            display_error(f"Error in response_async: {str(error)}")
-            raise
+            raise
+
+    def _generate_tool_definition(self, function_name: str) -> Optional[Dict]:
+        """Generate a tool definition from a function name."""
+        logging.debug(f"Attempting to generate tool definition for: {function_name}")
+        
+        # First try to get the tool definition if it exists
+        tool_def_name = f"{function_name}_definition"
+        tool_def = globals().get(tool_def_name)
+        logging.debug(f"Looking for {tool_def_name} in globals: {tool_def is not None}")
+        
+        if not tool_def:
+            import __main__
+            tool_def = getattr(__main__, tool_def_name, None)
+            logging.debug(f"Looking for {tool_def_name} in __main__: {tool_def is not None}")
+        
+        if tool_def:
+            logging.debug(f"Found tool definition: {tool_def}")
+            return tool_def
+
+        # Try to find the function
+        func = globals().get(function_name)
+        logging.debug(f"Looking for {function_name} in globals: {func is not None}")
+        
+        if not func:
+            import __main__
+            func = getattr(__main__, function_name, None)
+            logging.debug(f"Looking for {function_name} in __main__: {func is not None}")
+        
+        if not func or not callable(func):
+            logging.debug(f"Function {function_name} not found or not callable")
+            return None
+
+        import inspect
+        # Handle Langchain and CrewAI tools
+        if inspect.isclass(func) and hasattr(func, 'run') and not hasattr(func, '_run'):
+            original_func = func
+            func = func.run
+            function_name = original_func.__name__
+        elif inspect.isclass(func) and hasattr(func, '_run'):
+            original_func = func
+            func = func._run
+            function_name = original_func.__name__
+
+        sig = inspect.signature(func)
+        logging.debug(f"Function signature: {sig}")
+        
+        # Skip self, *args, **kwargs
+        parameters_list = []
+        for name, param in sig.parameters.items():
+            if name == "self":
+                continue
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+                continue
+            parameters_list.append((name, param))
+
+        parameters = {
+            "type": "object",
+            "properties": {},
+            "required": []
+        }
+        
+        # Parse docstring for parameter descriptions
+        docstring = inspect.getdoc(func)
+        logging.debug(f"Function docstring: {docstring}")
+        
+        param_descriptions = {}
+        if docstring:
+            import re
+            param_section = re.split(r'\s*Args:\s*', docstring)
+            logging.debug(f"Param section split: {param_section}")
+            if len(param_section) > 1:
+                param_lines = param_section[1].split('\n')
+                for line in param_lines:
+                    line = line.strip()
+                    if line and ':' in line:
+                        param_name, param_desc = line.split(':', 1)
+                        param_descriptions[param_name.strip()] = param_desc.strip()
+        
+        logging.debug(f"Parameter descriptions: {param_descriptions}")
+
+        for name, param in parameters_list:
+            param_type = "string"  # Default type
+            if param.annotation != inspect.Parameter.empty:
+                if param.annotation == int:
+                    param_type = "integer"
+                elif param.annotation == float:
+                    param_type = "number"
+                elif param.annotation == bool:
+                    param_type = "boolean"
+                elif param.annotation == list:
+                    param_type = "array"
+                elif param.annotation == dict:
+                    param_type = "object"
+            
+            parameters["properties"][name] = {
+                "type": param_type,
+                "description": param_descriptions.get(name, "Parameter description not available")
+            }
+            
+            if param.default == inspect.Parameter.empty:
+                parameters["required"].append(name)
+        
+        logging.debug(f"Generated parameters: {parameters}")
+        tool_def = {
+            "type": "function",
+            "function": {
+                "name": function_name,
+                "description": docstring.split('\n\n')[0] if docstring else "No description available",
+                "parameters": parameters
+            }
+        }
+        logging.debug(f"Generated tool definition: {tool_def}")
+        return tool_def
@@ -86,6 +86,7 @@ async def execute_callback(display_type: str, **kwargs):
 def _clean_display_content(content: str, max_length: int = 20000) -> str:
    """Helper function to clean and truncate content for display."""
    if not content or not str(content).strip():
+        logging.debug(f"Empty content received in _clean_display_content: {repr(content)}")
        return ""
        
    content = str(content)
@@ -174,11 +175,14 @@ def display_instruction(message: str, console=None, agent_name: str = None, agen
        console.print(Panel.fit(Text(message, style="bold blue"), title="Instruction", border_style="cyan"))

 def display_tool_call(message: str, console=None):
+    logging.debug(f"display_tool_call called with message: {repr(message)}")
    if not message or not message.strip():
+        logging.debug("Empty message in display_tool_call, returning early")
        return
    if console is None:
        console = Console()
    message = _clean_display_content(str(message))
+    logging.debug(f"Cleaned message in display_tool_call: {repr(message)}")
    
    # Execute callback if registered
    if 'tool_call' in sync_display_callbacks:
@@ -202,7 +206,8 @@ def display_error(message: str, console=None):

 def display_generating(content: str = "", start_time: Optional[float] = None):
    if not content or not str(content).strip():
-        return Panel("", title="", border_style="green")
+        logging.debug("Empty content in display_generating, returning early")
+        return None
    
    elapsed_str = ""
    if start_time is not None:
@@ -293,11 +298,14 @@ async def adisplay_instruction(message: str, console=None, agent_name: str = Non

 async def adisplay_tool_call(message: str, console=None):
    """Async version of display_tool_call."""
+    logging.debug(f"adisplay_tool_call called with message: {repr(message)}")
    if not message or not message.strip():
+        logging.debug("Empty message in adisplay_tool_call, returning early")
        return
    if console is None:
        console = Console()
    message = _clean_display_content(str(message))
+    logging.debug(f"Cleaned message in adisplay_tool_call: {repr(message)}")
    
    if 'tool_call' in async_display_callbacks:
        await async_display_callbacks['tool_call'](message=message)
@@ -321,7 +329,8 @@ async def adisplay_error(message: str, console=None):
 async def adisplay_generating(content: str = "", start_time: Optional[float] = None):
    """Async version of display_generating."""
    if not content or not str(content).strip():
-        return Panel("", title="", border_style="green")
+        logging.debug("Empty content in adisplay_generating, returning early")
+        return None
    
    elapsed_str = ""
    if start_time is not None:
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"

 [project]
 name = "praisonaiagents"
-version = "0.0.61"
+version = "0.0.65"
 description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
 authors = [
    { name="Mervin Praison" }
@@ -1868,7 +1868,7 @@ wheels = [

 [[package]]
 name = "praisonaiagents"
-version = "0.0.61"
+version = "0.0.65"
 source = { editable = "." }
 dependencies = [
    { name = "openai" },
@@ -3060,7 +3060,7 @@ wheels = [

 [[package]]
 name = "praisonai"
-version = "2.0.62"
+version = "2.0.80"
 source = { editable = "." }
 dependencies = [
    { name = "instructor" },
@@ -3197,7 +3197,7 @@ requires-dist = [
    { name = "plotly", marker = "extra == 'realtime'", specifier = ">=5.24.0" },
    { name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.0.7" },
    { name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.0.7" },
-    { name = "praisonaiagents", specifier = ">=0.0.61" },
+    { name = "praisonaiagents", specifier = ">=0.0.65" },
    { name = "pyautogen", marker = "extra == 'autogen'", specifier = ">=0.2.19" },
    { name = "pydantic", marker = "extra == 'chat'", specifier = "<=2.10.1" },
    { name = "pydantic", marker = "extra == 'code'", specifier = "<=2.10.1" },
@@ -3250,16 +3250,16 @@ wheels = [

 [[package]]
 name = "praisonaiagents"
-version = "0.0.61"
+version = "0.0.65"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "openai" },
    { name = "pydantic" },
    { name = "rich" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/50/be/09af17ea9eee5135c7b091bca77e1fc7a8ff732a6148497276ff5a06bd2e/praisonaiagents-0.0.61.tar.gz", hash = "sha256:9e178b99a1d2ffe23acd8413810e5d937d910a72fbd851e04d10b0bcb129b48a", size = 102057 }
+sdist = { url = "https://files.pythonhosted.org/packages/9e/58/6e07c4abf57d78be642f1517fe0692a003dfc253ad91a03c70521d246b5e/praisonaiagents-0.0.65.tar.gz", hash = "sha256:16627166f1b42bd8d21759fbd67015fff0b4730ca224b4a566cffe2b5bb770e8", size = 105128 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/12/1b/5656b79ae6c9f4c4e9e7b22db6800841c3bbab7fbd73c6c7f3381c936e64/praisonaiagents-0.0.61-py3-none-any.whl", hash = "sha256:835ec8a6fb1b2805a150db3d38bd4a1fe88333f078ee00190b102a3a4b2687a6", size = 121044 },
+    { url = "https://files.pythonhosted.org/packages/04/02/e4f2ef3c008a407d3b945256d1513bf0159989055e9462c691265be3d294/praisonaiagents-0.0.65-py3-none-any.whl", hash = "sha256:f53e9c3ddf31c575dd5f65315db70b35487adb0f9ae908596e74ed315c12afc5", size = 124306 },
 ]

 [[package]]
Author	SHA1	Message	Date
Mervin Praison	d278f2d2b2	Merge pull request #399 from MervinPraison/develop Bump version to 2.0.80 and update praisonaiagents dependency	2025-02-28 10:25:47 +00:00
MervinPraison	df4596a699	Bump version to 2.0.80 and update praisonaiagents dependency - Updated version to 2.0.80 across multiple files including praisonai.rb, pyproject.toml, uv.lock, Dockerfile, and docs - Upgraded praisonaiagents dependency from 0.0.64 to 0.0.65 - Enhanced agent start methods with optional return_dict parameter - Improved result handling in start and astart methods to return final response by default - Synchronized version updates across project configuration files	2025-02-28 10:25:32 +00:00
Mervin Praison	dd0f47ec81	Merge pull request #398 from MervinPraison/develop Bump version to 2.0.79 and update praisonaiagents dependency	2025-02-28 07:26:52 +00:00
MervinPraison	e53017a2fe	Bump version to 2.0.79 and update praisonaiagents dependency - Updated version to 2.0.79 across multiple files including praisonai.rb, pyproject.toml, uv.lock, Dockerfile, and docs - Upgraded praisonaiagents dependency from 0.0.63 to 0.0.64 - Updated example agent scripts for Anthropic and Gemini LLMs - Added parameter modification support in LLM initialization - Synchronized version updates across project configuration files	2025-02-28 07:26:22 +00:00
Mervin Praison	fe034ba8f4	Merge pull request #397 from MervinPraison/develop Enhance Agent Tool Execution Logging and Debugging	2025-02-28 07:05:13 +00:00
MervinPraison	063e5780d6	Enhance Agent Tool Execution Logging and Debugging - Added comprehensive debug logging for tool execution in LLM and Agent components - Improved display of tool call messages with more detailed information - Added logging for empty content and tool execution scenarios - Enhanced error handling and logging in display functions - Added new Anthropic LLM example script in src/praisonai-agents - Bumped version to 2.0.78 and updated praisonaiagents dependency to 0.0.63	2025-02-28 07:04:50 +00:00
Mervin Praison	904e21ca87	Merge pull request #396 from MervinPraison/develop v2.0.77	2025-02-27 09:12:33 +00:00
MervinPraison	526665ea11	Bump version to 2.0.77 and update praisonaiagents dependency - Updated version to 2.0.77 across multiple files including praisonai.rb, pyproject.toml, uv.lock, Dockerfile, and docs - Upgraded praisonaiagents dependency from 0.0.61 to 0.0.62 - Synchronized version updates across project configuration files	2025-02-27 09:11:18 +00:00
MervinPraison	c2e97d2432	adding litellm tools fix	2025-02-27 09:08:22 +00:00
MervinPraison	60be6a1f99	Add Comprehensive Debug Logging to LLM and Agent Components - Enhanced LLM and Agent classes with detailed debug logging capabilities - Added parameter and configuration logging for get_response and chat methods - Implemented time tracking for method execution in debug mode - Masked sensitive information like API keys in logs - Created new Wikipedia agent example script demonstrating agent initialization	2025-02-27 06:54:32 +00:00
Mervin Praison	0ba7d2e834	Merge pull request #385 from MervinPraison/develop Add Vision Model Upload Functionality and Enhance Training Workflow	2025-02-09 21:23:38 +00:00
MervinPraison	184b241395	Add Vision Model Upload Functionality and Enhance Training Workflow - Introduced new upload_vision.py script for comprehensive model upload capabilities - Extended train_vision.py with additional upload methods for Hugging Face and Ollama - Added support for saving merged models, pushing GGUF models, and creating Ollama models - Implemented flexible upload targeting with configurable options - Enhanced model preparation and deployment workflow for vision models	2025-02-09 21:23:24 +00:00
Mervin Praison	4354e76fe5	Merge pull request #383 from MervinPraison/develop Enhance Vision Model Training with Robust Dataset Conversion and Memo…	2025-02-09 02:46:01 +00:00
MervinPraison	4dd77efb75	Enhance Vision Model Training with Robust Dataset Conversion and Memory Optimization - Significantly improved train_vision.py with more robust dataset conversion process - Added detailed debugging and logging for dataset loading and conversion - Implemented advanced memory management techniques for vision model training - Enhanced tokenizer handling with fallback mechanisms - Updated version to 2.0.76 across all relevant files	2025-02-09 02:45:47 +00:00
Mervin Praison	11db6e18ba	Merge pull request #382 from MervinPraison/develop Add Vision Model Training Support and Update Unsloth Dependency	2025-02-09 00:31:40 +00:00
MervinPraison	0d6f698e6c	Add Vision Model Training Support and Update Unsloth Dependency - Enhanced CLI to dynamically select training script based on model type (vision vs. standard) - Added support for detecting vision models through flexible name matching - Updated Unsloth dependency to latest commit - Bumped version to 2.0.75 across all relevant files	2025-02-09 00:31:20 +00:00
Mervin Praison	b8574f78e2	Merge pull request #381 from MervinPraison/develop Bump version to 2.0.74 and refactor train.py template generation	2025-02-07 13:14:32 +00:00
MervinPraison	e631774a0f	Bump version to 2.0.74 and refactor train.py template generation - Updated version to 2.0.74 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Simplified training template generation in train.py - Removed complex stop tokens and simplified template structure - Dynamically built training arguments with improved configuration flexibility	2025-02-07 13:14:16 +00:00
Mervin Praison	df735aa748	Merge pull request #378 from MervinPraison/develop Enhance train.py with Configurable Training Parameters	2025-02-06 07:48:00 +00:00
MervinPraison	3bdf5b6c65	Enhance train.py with Configurable Training Parameters - Updated version to 2.0.73 across multiple files - Refactored TrainModel to use configurable training parameters with default fallbacks - Improved flexibility of training configuration by allowing custom settings via config - Dynamically set training arguments based on provided configuration	2025-02-06 07:47:40 +00:00
Mervin Praison	361fa67108	Merge pull request #377 from MervinPraison/develop Bump version to 2.0.72 and enhance CLI configuration in train and cli…	2025-02-06 07:26:59 +00:00
MervinPraison	c5abcf362a	Bump version to 2.0.72 and enhance CLI configuration in train and cli modules - Updated version to 2.0.72 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Refined CLI configuration handling in cli.py for config.yaml generation - Simplified train.py main function to always use provided config path - Improved configuration flexibility for training and model setup	2025-02-06 07:26:41 +00:00
Mervin Praison	7ed2d99a11	Merge pull request #376 from MervinPraison/develop Enhance train.py CLI configuration handling and bump version to 2.0.71	2025-02-06 07:19:11 +00:00
MervinPraison	81307da7c8	Enhance train.py CLI configuration handling and bump version to 2.0.71 - Updated version to 2.0.71 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Modified train.py to use existing config.yaml when no CLI arguments are provided - Improved CLI configuration flexibility for training command	2025-02-06 07:18:56 +00:00
Mervin Praison	fb75153b5c	Merge pull request #373 from MervinPraison/develop Bump version to 2.0.70 and update configuration handling	2025-02-05 06:28:07 +00:00
MervinPraison	cddcfcf170	Bump version to 2.0.70 and update configuration handling - Updated version to 2.0.70 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Enhanced CLI configuration handling in cli.py to load existing config when not creating a new one - Minor improvements to configuration management and deployment scripts	2025-02-05 06:27:50 +00:00
Mervin Praison	398b0473e0	Merge pull request #372 from MervinPraison/develop Update train.mdx with Wandb configuration and requirements	2025-02-05 05:07:23 +00:00
MervinPraison	c5d260782a	Update train.mdx with Wandb configuration and requirements - Added detailed requirements for training, including Huggingface token and model details - Included Wandb login and configuration instructions - Updated documentation with environment variables for Wandb logging	2025-02-05 05:07:08 +00:00
Mervin Praison	116ac5a64d	Merge pull request #370 from MervinPraison/develop Add Wandb Support for Training and Bump Version to 2.0.69	2025-02-04 11:18:42 +00:00
MervinPraison	abf6e2e8d1	Add Wandb Support for Training and Bump Version to 2.0.69 - Updated version to 2.0.69 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Enhanced train.py with optional Wandb logging support - Added environment variable-based Wandb configuration - Implemented dynamic Wandb reporting and run naming	2025-02-04 11:18:20 +00:00
Mervin Praison	cb2b3617fa	Merge pull request #369 from MervinPraison/develop Bump version to 2.0.68 and enhance Ollama model template generation	2025-02-04 11:10:50 +00:00
MervinPraison	6de67c7d0d	Bump version to 2.0.68 and enhance Ollama model template generation - Updated version across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Significantly improved train.py's prepare_modelfile_content() method - Added comprehensive template and stop token mapping for multiple model architectures - Implemented dynamic template selection based on model name - Enhanced Ollama model file generation with more flexible configuration	2025-02-04 11:10:26 +00:00
Mervin Praison	5699d90894	Merge pull request #368 from MervinPraison/develop Enhance train.py with Dynamic Stop Token Generation for Ollama Models	2025-02-04 10:33:19 +00:00
MervinPraison	2982e8ea9c	Enhance train.py with Dynamic Stop Token Generation for Ollama Models - Added intelligent stop token inference for various model architectures in train.py - Implemented dynamic stop token selection based on model name - Supported models include Phi, LLaVA, Mistral, Qwen, DeepSeek, and default models - Updated Modelfile generation to include model-specific stop tokens - Improved Ollama model preparation flexibility	2025-02-04 10:33:05 +00:00
Mervin Praison	9dd2146342	Merge pull request #367 from MervinPraison/develop Bump version to 2.0.66 and remove TRAIN_AVAILABLE flag	2025-02-04 09:31:20 +00:00
MervinPraison	f79184a6ee	Bump version to 2.0.66 and remove TRAIN_AVAILABLE flag - Updated version across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, deploy.py, docs) - Removed TRAIN_AVAILABLE flag from cli.py - Simplified CLI configuration by removing optional training feature check	2025-02-04 09:31:07 +00:00
Mervin Praison	0985e12ad4	Merge pull request #366 from MervinPraison/develop Bump version to 2.0.65 and update deployment configurations	2025-02-04 08:05:06 +00:00
MervinPraison	0fbb7972e1	Bump version to 2.0.65 and update deployment configurations - Updated version across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, deploy.py, docs) - Enhanced Conda environment setup script with Ollama installation check - Minor improvements to deployment and setup scripts	2025-02-04 08:04:44 +00:00
Mervin Praison	5744ae853e	Merge pull request #365 from MervinPraison/develop Develop	2025-02-04 07:36:59 +00:00
MervinPraison	bd6e3da4ff	Bump version to 2.0.64 and update dependencies - Updated version across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, deploy.py, docs) - Added TRAIN_AVAILABLE flag in cli.py for optional training feature - Updated Conda environment setup script with latest dependencies - Minor improvements to training and deployment configurations	2025-02-04 07:36:43 +00:00
MervinPraison	f2f51d9e8b	Refactor train.py with Enhanced Dataset Processing and Tokenization - Completely rewrote training script with robust dataset handling - Added support for ShareGPT and Alpaca-style datasets - Implemented advanced formatting and tokenization functions - Enhanced debug logging and error handling - Simplified model preparation and training workflow - Updated chat template and tokenization strategies - Improved inference and model saving methods	2025-02-03 19:38:57 +00:00
Mervin Praison	4a1098603d	Merge pull request #364 from MervinPraison/develop cli docs	2025-02-03 07:41:14 +00:00
MervinPraison	5468e35bc1	cli docs	2025-02-03 07:40:46 +00:00
Mervin Praison	b468267a87	Merge pull request #362 from MervinPraison/develop Develop	2025-02-02 06:45:30 +00:00
MervinPraison	48141bbd9e	update readme	2025-02-02 06:45:12 +00:00
MervinPraison	b3cdf53d07	adding demo	2025-02-02 06:43:16 +00:00
Mervin Praison	aa594ddc32	Merge pull request #361 from MervinPraison/develop update read me 2	2025-02-02 06:30:57 +00:00
MervinPraison	5f330e8fb2	update read me 2	2025-02-02 06:30:37 +00:00
Mervin Praison	ff94e128de	Merge pull request #360 from MervinPraison/develop update read me	2025-02-02 06:28:35 +00:00
MervinPraison	854a016fc5	update read me	2025-02-02 06:28:08 +00:00
Mervin Praison	12e8ed7b77	Merge pull request #359 from MervinPraison/develop Update CLI Documentation with Latest LLM Model Option	2025-02-01 21:51:34 +00:00
MervinPraison	13c5f5537e	Update CLI Documentation with Latest LLM Model Option - Updated CLI documentation to reflect the latest LLM model option - Changed example command to use `gpt-4o-mini` instead of `gpt-4-mini` - Minor documentation refinement for CLI feature	2025-02-01 21:51:20 +00:00
Mervin Praison	e7c1614ac8	Merge pull request #358 from MervinPraison/develop Add CLI Documentation and Enhance CLI Functionality	2025-02-01 00:20:09 +00:00
MervinPraison	7cc6db2a34	Add CLI Documentation and Enhance CLI Functionality - Added comprehensive CLI documentation in docs/features/cli.mdx - Enhanced CLI functionality with new command-line options - Implemented `--llm` flag to specify LLM models for direct prompts - Updated CLI parsing to support flexible LLM model selection - Improved direct prompt handling across different AI frameworks - Bumped package version to 2.0.63	2025-02-01 00:19:44 +00:00