62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
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 |