cyclop/tools/standalone_llm_tool.py

from .base_tool import BaseTool
import os
import json
import logging
from openai import OpenAI

class StandaloneLLMTool(BaseTool):
    def __init__(self):
        self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    def clear(self):
        pass

    def get_functions(self):
        return [
            {
                "type": "function",
                    "function": {
                    "name": "call_external_llm",
                    "description": "Call an external language model",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "prompt": {
                                "type": "string",
                                "description": "The prompt you are providing"
                            },
                            "model": {
                                "type": "string",
                                "description": "The model to use for generating the detailed instructions. Use mini for most coding tasks, preview when needing sophisticated reasoning",
                                "enum": ["mini", "max"],
                                "default": "o1-mini"
                            },
                            "max_tokens": {
                                "type": "integer",
                                "description": "The maximum number of tokens to use for generating the detailed instructions. Default is 16384.",
                            }
                        },
                        "required": ["prompt"]
                    }
                },
                "_tags": ["llm", "external"]
            }
        ]

    def execute(self, function_name, **kwargs):
        if function_name == "call_external_llm":
            return self.call_external_llm(kwargs.get("prompt"), kwargs.get("model"), kwargs.get("max_tokens"))
        else:
            error_message = f"Unknown function: {function_name}"
            logging.error(error_message)

    def call_external_llm(self, prompt, model="o1-mini", max_tokens=16384):
        logging.info(f"Calling external model: {model}")
        response = self.client.completions.create(
            model=model,
            prompt=prompt,
            max_tokens=max_tokens
        )
        token_amount = response.summary["total_tokens"]
        logging.info("Response generated, {token_amount} tokens used.")
        return response.choices[0].text